{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vRWRs6lypGD5"
      },
      "source": [
        "# License\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at:\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcwyNrywpM_R"
      },
      "source": [
        "# Instructions\n",
        "\n",
        "This Notebook allows to reproduce the experiments reported in the publication titled:\n",
        "\n",
        "\"[*A Multi-Agent Framework for the Asynchronous and Collaborative Extension of Multitask ML Systems*](https://arxiv.org/abs/2209.14745)\" (2022)\n",
        "\n",
        "---\n",
        "To start an experiment:\n",
        "---\n",
        "\n",
        "1. Set `EXPERIMENT_NAME` in the configurations below to a name of choice. This will be used also as name of the folder storing the state of the µNet system generated by the experiment (system state folder).\n",
        "\n",
        "1. By default the system state is written in a temporary folder stored in the memory of the current Virtual Machine (VM). This temporary folder is deleted every time the VM is stopped or restarted.\n",
        "It is possible to save the system state folder on your Google Drive by activating the\n",
        "`SAVE_SYSTEM_STATE_ON_GOOGLE_DRIVE` option. In this case you will be prompted for access approval and the system state folder will be saved in a folder named `\"munet_experiments\"` under your Google Drive root folder. It is also possible to store the state into a Google Drive folder shared with multiple users by creating a link to the shared folder into your Google Drive and then setting `EXPERIMENTS_ROOT_DIR` (below) to the path of the linked shared folder. \n",
        "\n",
        "1. `AGENT` can be set to: a) `VitT3` to use a ViT-Tiny root model capped to 3 layers and apply the experiment configuration used for the experiments on the *Multitask Character Classifiaction Benchmark*.\n",
        "b) `VitB` to use a ViT Base root model and apply the experiment configuration used for the experiments on the *Visual Domain Decathlong Benchmark*.\n",
        "c) `Vit` to use a ViT-Large root model and apply the large-scale continual-learning experiment configuration.\n",
        "\n",
        "1. Set `TASK_NAME` to the task id string.\n",
        "For example a task id from the [Tensorflow Datasets image classification catalog](https://www.tensorflow.org/datasets/catalog/beans).\n",
        "Refer to `TFDS_IMAGE_CLASSIFCATON_DATASETS` and `VTAB_TASKS` (below) for lists of tasks ids that have been tested with the current code. Note that some tasks require manual download, refer to the corresponding catalogue page for instructions. **WARNING**: The system state needs to be populated with at least one rooot model before running an agent training on any task. To generate the root model set `TASK_NAME` to either `\"root_model/checkpoint\"` or `\"root_model/random_init\"` for respectively loading a pretrained root model or generating a randomly initialized one.\n",
        "\n",
        "1. To start the experiment, select \"Connect to a hosted runtime\" from the dropdown menu on the top right, and then select \"Run all\" from the \"Runtime\" menu.\n",
        "\n",
        "---\n",
        "During the experiment execution:\n",
        "---\n",
        "\n",
        "1. The logging output is printed after the last cell of this Colab.\n",
        "\n",
        "1. The system state folder is populated with a subfolder for each agent.\n",
        "The name of each agent folder is prefixed with the agent name and sufixed with the task name.\n",
        "Each agent directory is populated with incremental state subfolders  containing the sharded state of the architectures and parameters generated by the agent.\n",
        "\n",
        "1. Agents can be started asyncronously and run in parallel in varying quantities.\n",
        "It is possible to resume an interrupted agent training by restarting the execution with the same configuration.\n",
        "It is possible to further continue a complented training by increasing the `config.num_cycles_max` set in the configurations below.\n",
        "\n",
        "1. To achieve a multi-agent execution, multiple Colabs need to be run in parallel, each set to the same configuration but different `TASK_NAME`.\n",
        "\n",
        "1. To achieve heterogeneous hardware execution, parallel Colab Notebooks can be connected to a runtime of different types.\n",
        "It is possible to switch between CPU, GPU and TPU by selecting `Change runtime type` in the `Resources` tab in this Colab Notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M93tll7z29rX"
      },
      "outputs": [],
      "source": [
        "# @title Agent parameters\n",
        "EXPERIMENT_NAME = \"munet_test\"  # @param { type: \"string\", isTemplate: true }\n",
        "SAVE_SYSTEM_STATE_ON_GOOGLE_DRIVE = False  # @param { type: \"boolean\", isTemplate: true }\n",
        "AGENT = \"VitT3\" # @param [\"VitT3\", \"VitB\", \"Vit\"] { type: \"string\", isTemplate: true }\n",
        "# Set TASK_NAME to \"root_model/checkpoint\" or \"root_model/random_init\" to initalize the population.\n",
        "TASK_NAME = \"root_model/checkpoint\"  # @param { type: \"string\", isTemplate: true }\n",
        "assert TASK_NAME\n",
        "\n",
        "# Set to False to disable autotune.\n",
        "AUTO_TUNE = True\n",
        "# Allows to override the default scale factor = 1 (no cost penalties)\n",
        "SCALE_FACTOR = None\n",
        "if SCALE_FACTOR is not None:\n",
        "  assert 0 \u003c SCALE_FACTOR \u003c= 1\n",
        "# Print debug statements.\n",
        "VERBOSE = False\n",
        "# Skip intermediate state save if last state was written within this time range.\n",
        "SKIP_INTERMEDIATE_STATE_SECS = 10 * 60  # 10 minutes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3E9UilpdXABD"
      },
      "outputs": [],
      "source": [
        "if SAVE_SYSTEM_STATE_ON_GOOGLE_DRIVE:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/gdrive')\n",
        "  EXPERIMENTS_ROOT_DIR = \"/content/gdrive/My Drive/munet_experiments/\"\n",
        "  print(\"Saving system state in Google Drive.\")\n",
        "else:\n",
        "  EXPERIMENTS_ROOT_DIR = \"/tmp/\"\n",
        "  print(\"WARNING: Saving system state in VM, changes will be lost after reboot!!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "34veUvnzyknm"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade -q pip jax jaxlib\n",
        "!pip install --upgrade -q git+https://github.com/google/flax.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DF968cFFymJS"
      },
      "outputs": [],
      "source": [
        "!pip install -q ml_collections"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M3K7AXG03VL-"
      },
      "outputs": [],
      "source": [
        "!pip install -q tensorflow_addons"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PW6LrPB0yoci"
      },
      "outputs": [],
      "source": [
        "![ -d task_adaptation ] || git clone --depth=1 https://github.com/google-research/task_adaptation\n",
        "![ -d vision_transformer ] || git clone --depth=1 https://github.com/google-research/vision_transformer\n",
        "import sys\n",
        "if './task_adaptation' not in sys.path:\n",
        "  sys.path.append('./task_adaptation')\n",
        "if './vision_transformer' not in sys.path:\n",
        "  sys.path.append('./vision_transformer')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-fOzt8Boyq97"
      },
      "outputs": [],
      "source": [
        "import jax.tools.colab_tpu\n",
        "try:\n",
        "  jax.tools.colab_tpu.setup_tpu()\n",
        "except:\n",
        "  pass  # Not a Tpu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NLPC-HkUyqZp"
      },
      "outputs": [],
      "source": [
        "import copy\n",
        "import datetime\n",
        "import gc\n",
        "import inspect\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import json\n",
        "import math\n",
        "import matplotlib\n",
        "import numpy as np\n",
        "import os\n",
        "import optax\n",
        "import pandas as pd\n",
        "import random\n",
        "import re\n",
        "import time\n",
        "from collections import defaultdict\n",
        "from functools import partial\n",
        "from matplotlib import pyplot as plt\n",
        "from threading import Thread, Lock\n",
        "from typing import Optional"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FAl7kkIf3g6L"
      },
      "outputs": [],
      "source": [
        "import flax\n",
        "import flax.linen as nn\n",
        "from flax.training import checkpoints as flax_checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U3dGgwM1yudH"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.io.gfile as gfile\n",
        "import tensorflow_datasets as tfds\n",
        "tf.compat.v1.enable_eager_execution()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Mx2ymKNyv9r"
      },
      "outputs": [],
      "source": [
        "from ml_collections import ConfigDict, FrozenConfigDict\n",
        "from vision_transformer.vit_jax import input_pipeline\n",
        "from vision_transformer.vit_jax import checkpoint\n",
        "from vision_transformer.vit_jax.configs import models as models_config  # Model configurations.\n",
        "from vision_transformer.vit_jax import models_vit as models # Actual model code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f8l2kxTU3vFA"
      },
      "outputs": [],
      "source": [
        "import task_adaptation.registry as task_adapt_registry\n",
        "import task_adaptation.data.caltech\n",
        "import task_adaptation.data.cifar\n",
        "import task_adaptation.data.dtd\n",
        "import task_adaptation.data.oxford_flowers102\n",
        "import task_adaptation.data.oxford_iiit_pet\n",
        "import task_adaptation.data.sun397\n",
        "import task_adaptation.data.svhn\n",
        "import task_adaptation.data.patch_camelyon\n",
        "import task_adaptation.data.eurosat\n",
        "import task_adaptation.data.resisc45\n",
        "import task_adaptation.data.diabetic_retinopathy\n",
        "import task_adaptation.data.clevr\n",
        "import task_adaptation.data.dmlab\n",
        "import task_adaptation.data.dsprites\n",
        "import task_adaptation.data.kitti\n",
        "import task_adaptation.data.smallnorb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SHBWj0JmpWDX"
      },
      "outputs": [],
      "source": [
        "# Ref Tfds catalog: https://www.tensorflow.org/datasets/catalog/beans\n",
        "TFDS_IMAGE_CLASSIFCATON_DATASETS = set([\n",
        "    \"beans\",\n",
        "    \"binary_alpha_digits\",\n",
        "    \"caltech_birds2010\",\n",
        "    \"caltech_birds2011\",\n",
        "    \"cars196\",\n",
        "    \"cassava\",\n",
        "    \"cats_vs_dogs\",\n",
        "    \"cifar10\",\n",
        "    \"cifar100\",\n",
        "    \"citrus_leaves\",\n",
        "    \"cmaterdb/bangla\",\n",
        "    \"cmaterdb/devanagari\",\n",
        "    \"cmaterdb/telugu\",\n",
        "    \"colorectal_histology\",\n",
        "    \"controlled_noisy_web_labels/mini_imagenet_red\",\n",
        "    \"controlled_noisy_web_labels/mini_imagenet_blue\",\n",
        "    \"curated_breast_imaging_ddsm/patches\",\n",
        "    \"cycle_gan/apple2orange\",\n",
        "    \"cycle_gan/summer2winter_yosemite\",\n",
        "    \"cycle_gan/horse2zebra\",\n",
        "    \"cycle_gan/monet2photo\",\n",
        "    \"cycle_gan/cezanne2photo\",\n",
        "    \"cycle_gan/ukiyoe2photo\",\n",
        "    \"cycle_gan/vangogh2photo\",\n",
        "    \"cycle_gan/maps\",\n",
        "    \"cycle_gan/cityscapes\",\n",
        "    \"cycle_gan/facades\",\n",
        "    \"cycle_gan/iphone2dslr_flower\",\n",
        "    \"deep_weeds\",\n",
        "    \"domainnet/real\",\n",
        "    \"domainnet/painting\",\n",
        "    \"domainnet/clipart\",\n",
        "    \"domainnet/quickdraw\",\n",
        "    \"domainnet/infograph\",\n",
        "    \"domainnet/sketch\",\n",
        "    \"emnist/balanced\",\n",
        "    \"emnist/byclass\",\n",
        "    \"emnist/bymerge\",\n",
        "    \"emnist/digits\",\n",
        "    \"emnist/letters\",\n",
        "    \"emnist/mnist\",\n",
        "    \"fashion_mnist\",\n",
        "    \"food101\",\n",
        "    \"horses_or_humans\",\n",
        "    \"i_naturalist2017\",\n",
        "    \"i_naturalist2018\",\n",
        "    \"imagenet2012\",\n",
        "    \"imagenet_a\",\n",
        "    \"imagenet_lt\",\n",
        "    \"imagenet_r\",\n",
        "    \"imagenet_sketch\",\n",
        "    \"imagenette\",\n",
        "    \"imagewang\",\n",
        "    \"kmnist\",\n",
        "    \"malaria\",\n",
        "    \"mnist\",\n",
        "    \"omniglot\",\n",
        "    \"pet_finder\",\n",
        "    \"places365_small\",\n",
        "    \"plant_village\",\n",
        "    \"plantae_k\",\n",
        "    \"quickdraw_bitmap\",\n",
        "    \"rock_paper_scissors\",\n",
        "    \"siscore/rotation\",\n",
        "    \"siscore/size\",\n",
        "    \"siscore/location\",\n",
        "    \"stanford_dogs\",\n",
        "    \"stanford_online_products\",\n",
        "    \"stl10\",\n",
        "    \"tf_flowers\",\n",
        "    \"uc_merced\",\n",
        "    \"visual_domain_decathlon/aircraft\",\n",
        "    \"visual_domain_decathlon/cifar100\",\n",
        "    \"visual_domain_decathlon/daimlerpedcls\",\n",
        "    \"visual_domain_decathlon/dtd\",\n",
        "    \"visual_domain_decathlon/gtsrb\",\n",
        "    \"visual_domain_decathlon/imagenet12\",\n",
        "    \"visual_domain_decathlon/omniglot\",\n",
        "    \"visual_domain_decathlon/svhn\",\n",
        "    \"visual_domain_decathlon/ucf101\",\n",
        "    \"visual_domain_decathlon/vgg-flowers\",\n",
        "    ])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XI7xcfWLXx-G"
      },
      "outputs": [],
      "source": [
        "# Append suffix \"/1k\" to get the 1k version of each task.\n",
        "VTAB_TASKS = [\n",
        "              \"caltech101\",\n",
        "              # cifar100/10 were already added with slightly different val split\n",
        "              # but same test set. So here is added only the 1k versions.\n",
        "              \"cifar100/1k\",\n",
        "              \"cifar10/1k\",\n",
        "              \"dtd\",\n",
        "              \"oxford_flowers102\",\n",
        "              \"oxford_iiit_pet\",\n",
        "              \"sun397\",\n",
        "              \"svhn_cropped\",\n",
        "                ###\n",
        "              \"patch_camelyon\",\n",
        "              \"eurosat\",\n",
        "              \"resisc45\",\n",
        "              \"diabetic_retinopathy_detection/btgraham-300\",\n",
        "                ###\n",
        "              \"clevr/count_cylinders\",  # Not in results table.\n",
        "              \"clevr/count_all\",  # Clevr-Count\n",
        "              \"clevr/closest_object_distance\",  # Clevr-Dist\n",
        "              \"dmlab\",\n",
        "              \"dsprites/label_x_position\",  # dSpr-Loc\n",
        "              \"dsprites/label_orientation\",  # dSpr-Ori\n",
        "              \"kitti/closest_object_distance\",  # Not in results table.\n",
        "              \"kitti/count_vehicles\",  # Not in results table.\n",
        "              \"kitti/closest_vehicle_distance\",  # Kitti-dist\n",
        "              \"smallnorb/label_category\",  # Not in results table.\n",
        "              \"smallnorb/label_lighting\",  # Not in results table.\n",
        "              \"smallnorb/label_azimuth\",  # Azim\n",
        "              \"smallnorb/label_elevation\",  # Elev\n",
        "              ]\n",
        "\n",
        "for tn in VTAB_TASKS:\n",
        "  assert tn not in TFDS_IMAGE_CLASSIFCATON_DATASETS, tn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "deuT4iMO4Nhf"
      },
      "outputs": [],
      "source": [
        "TFDS_BUILDERS_CACHE = {}\n",
        "\n",
        "def get_tfds_builder(tfds_name):\n",
        "  global TFDS_BUILDERS_CACHE\n",
        "  if tfds_name not in TFDS_BUILDERS_CACHE:\n",
        "    TFDS_BUILDERS_CACHE[tfds_name] = tfds.builder(tfds_name)\n",
        "    TFDS_BUILDERS_CACHE[tfds_name].download_and_prepare()\n",
        "  return TFDS_BUILDERS_CACHE[tfds_name]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ExYTVeamDcPy"
      },
      "source": [
        "# ViT"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jTOBlYBpWNMl"
      },
      "outputs": [],
      "source": [
        "def ids_str2ints(ids_str):\n",
        "  return [int(v) for v in str(ids_str).split(\"_\")] if ids_str else []\n",
        "def ids_ints2str(ids_ints):\n",
        "  return \"_\".join([str(v) for v in sorted(ids_ints)])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e88M0y8YVZrC"
      },
      "outputs": [],
      "source": [
        "AddPositionEmbs = models.AddPositionEmbs\n",
        "Encoder1DBlock = models.Encoder1DBlock\n",
        "VisionTransformer = models.VisionTransformer\n",
        "\n",
        "class ResidualAdapter(nn.Module):\n",
        "  adapter_dim: int\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    hidden_dim = x.shape[-1]\n",
        "    y = nn.LayerNorm()(x)\n",
        "    y = nn.Dense(self.adapter_dim)(y)\n",
        "    y = nn.gelu(y)\n",
        "    # Zero Initialization so that added adapter does not change the representation.\n",
        "    y = nn.Dense(hidden_dim, kernel_init=jax.nn.initializers.zeros)(y)\n",
        "    return x + y  # Residual.\n",
        "\n",
        "# Modified from vision_transformer/vit_jax/models Encoder to add residual adapters.\n",
        "class Encoder(nn.Module):\n",
        "  num_layers: int\n",
        "  mlp_dim: int\n",
        "  num_heads: int\n",
        "  adapter_layers: str  # \u003cMOD\n",
        "  adapter_dim: int  # MOD\u003e\n",
        "  dropout_rate: float = 0.1\n",
        "  attention_dropout_rate: float = 0.1\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, inputs, *, train):\n",
        "    assert inputs.ndim == 3  # (batch, len, emb)\n",
        "\n",
        "    x = AddPositionEmbs(\n",
        "        posemb_init=nn.initializers.normal(stddev=0.02),  # from BERT.\n",
        "        name=\"posembed_input\")(\n",
        "            inputs)\n",
        "\n",
        "    # Input Encoder\n",
        "    adapter_layers_ids = ids_str2ints(self.adapter_layers)  # \u003cMOD\u003e\n",
        "    for lyr in range(self.num_layers):\n",
        "      if lyr in adapter_layers_ids:  # \u003cMOD\n",
        "        x = ResidualAdapter(\n",
        "            adapter_dim=self.adapter_dim,\n",
        "            name=f\"residual_adapter_{lyr}\"\n",
        "            )(x)  # MOD\u003e\n",
        "      x = Encoder1DBlock(\n",
        "          mlp_dim=self.mlp_dim,\n",
        "          dropout_rate=self.dropout_rate,\n",
        "          attention_dropout_rate=self.attention_dropout_rate,\n",
        "          name=f\"encoderblock_{lyr}\",\n",
        "          num_heads=self.num_heads)(x,\n",
        "                                    deterministic=True)  # Disable dropout.\n",
        "    encoded = nn.LayerNorm(name=\"encoder_norm\")(x)\n",
        "    return encoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vvZ_4-kJ9Pt3"
      },
      "outputs": [],
      "source": [
        "def get_vit_filename(query):\n",
        "  df = checkpoint.get_augreg_df()\n",
        "  res = df.query(query).filename.unique()\n",
        "  assert len(res) == 1\n",
        "  return res[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0lvqd47g9ZsW"
      },
      "outputs": [],
      "source": [
        "VIT_CONFIG_CACHE = {}\n",
        "\n",
        "def get_vit_config(query):\n",
        "  if query not in VIT_CONFIG_CACHE:\n",
        "    filename = get_vit_filename(query)\n",
        "    config = models_config.AUGREG_CONFIGS[filename.split(\"-\")[0]].copy_and_resolve_references()\n",
        "    # Overwrite with custom Encoder.\n",
        "    config.unlock()\n",
        "    config.encoder = Encoder\n",
        "    config.transformer.adapter_layers = \"\"\n",
        "    config.transformer.adapter_dim = -1\n",
        "    # Disable dropout.\n",
        "    config.transformer.dropout_rate = 0.0\n",
        "    config.transformer.attention_dropout_rate = 0.0\n",
        "    config.lock()\n",
        "    VIT_CONFIG_CACHE[query] = config\n",
        "  return VIT_CONFIG_CACHE[query].copy_and_resolve_references()\n",
        "\n",
        "def get_max_num_layers(query):\n",
        "  config = get_vit_config(query)\n",
        "  return config.transformer.num_layers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lyX6WjaUDWOi"
      },
      "source": [
        "# Agents"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9I56DcbI79M2"
      },
      "source": [
        "## Agents configs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2AwvZWRxnao5"
      },
      "outputs": [],
      "source": [
        "DATASET_HPARAMS_KEYS_PRERFIX = \"ds_\"\n",
        "OPTIMIZER_HPARAMS_KEYS_PRERFIX = \"opt_\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ks7L_IFVng3r"
      },
      "outputs": [],
      "source": [
        "def get_config_ti3_chars():\n",
        "  config = ConfigDict()\n",
        "  config.num_train_examples_between_validations_max = 51200  # 100 batches.\n",
        "  config.num_validations_per_path_training = 4\n",
        "  config.num_validation_examples_max = 5120  # 10 batches.\n",
        "  config.batch_size = 512\n",
        "  config.num_cycles_max = 5\n",
        "  config.num_samples_per_cycle = 4*8\n",
        "  # Force finetune last layer norm that technically is part of the head.\n",
        "  config.force_mutations = [\"clone:encoder_norm\"]\n",
        "  config.policy_class = \"PPDecay\"\n",
        "  config.policy_kwargs = {}\n",
        "  config.scorer_class = \"ScorerDecay\"\n",
        "  config.scorer_kwargs = dict(\n",
        "      scale_factor=1.0,\n",
        "      num_params=1_484_162,  # Params with Ti/16 3 layers.\n",
        "      flops=15_000_000,  # Flops with Ti/16 3 layers and image size 32.\n",
        "      )\n",
        "\n",
        "  config.vit_checkpoint_query = 'name==\"Ti/16\" and ds==\"i21k\" and aug==\"light1\" and wd==0.1 and sd==0.0'\n",
        "\n",
        "  config.models_default_hparams = {\n",
        "      \"_mu_\": 0.1,\n",
        "      # Default num_classes has no effect since it is always overwritten or used\n",
        "      # for rand init models whose head is always replaced.\n",
        "      \"num_classes\": 1,\n",
        "      \"mutate_adapters\": True,\n",
        "      # Set to ids_ints2str(range(max_num_layers)) to activate all adapters.\n",
        "      \"adapter_layers\": \"\",\n",
        "      \"num_layers\": 3,\n",
        "      \"adapter_dim\": 32,\n",
        "      \"opt_lr\": 0.01,\n",
        "      \"opt_lr_schedule\": \"cosine\",\n",
        "      \"opt_lr_warmup_ratio\": 0.1,\n",
        "      \"opt_momentum\": 0.9,\n",
        "      \"opt_nesterov\": False,\n",
        "      \"ds_image_size\": 32,\n",
        "      \"ds_area_range_min\": 1.0,\n",
        "      \"ds_aspect_ratio_range_min\": 1.0,\n",
        "      \"ds_flip_left_right\": False,\n",
        "      \"ds_brightness_delta\": 0.0,\n",
        "      \"ds_contrast_delta\": 0.0,\n",
        "      \"ds_saturation_delta\": 0.0,\n",
        "      \"ds_hue_delta\": 0.0,\n",
        "      \"ds_quality_delta\": 0.0,\n",
        "  }\n",
        "  config.models_mutation_ranges = {}\n",
        "  return config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_ToPBiVCnlK4"
      },
      "outputs": [],
      "source": [
        "def get_config_base_deca():\n",
        "  config = ConfigDict()\n",
        "  config.num_train_examples_between_validations_max = 51200\n",
        "  config.num_validations_per_path_training = 4\n",
        "  config.num_validation_examples_max = 5120\n",
        "  config.batch_size = 128\n",
        "  config.num_cycles_max = 10  # Set to 30 for convergence.\n",
        "  config.num_samples_per_cycle = 4*8\n",
        "  config.force_mutations = [\"clone:encoder_norm\"]\n",
        "  config.policy_class = \"PPDecay\"\n",
        "  config.policy_kwargs = {}\n",
        "  config.scorer_class = \"ScorerDecay\"\n",
        "  config.scorer_kwargs = dict(\n",
        "      scale_factor=1.0,\n",
        "      num_params=85_652_738,\n",
        "      flops=100_000_000_000,\n",
        "      )\n",
        "  config.vit_checkpoint_query = 'name==\"B/16\" and ds==\"i21k\" and aug==\"medium1\" and wd==0.1 and sd==0'\n",
        "\n",
        "  max_num_layers = get_max_num_layers(config.vit_checkpoint_query)\n",
        "  config.models_default_hparams = {\n",
        "      \"_mu_\": 0.1,\n",
        "      \"num_classes\": 1,\n",
        "      \"mutate_adapters\": True,\n",
        "      \"adapter_layers\": \"\",\n",
        "      \"num_layers\": max_num_layers,\n",
        "      \"adapter_dim\": 32,\n",
        "      \"opt_lr\": 0.01,\n",
        "      \"opt_lr_schedule\": \"cosine\",\n",
        "      \"opt_lr_warmup_ratio\": 0.1,\n",
        "      \"opt_momentum\": 0.9,\n",
        "      \"opt_nesterov\": False,\n",
        "      \"ds_image_size\": 80,\n",
        "      \"ds_area_range_min\": 1.0,\n",
        "      \"ds_aspect_ratio_range_min\": 1.0,\n",
        "      \"ds_flip_left_right\": False,\n",
        "      \"ds_brightness_delta\": 0.0,\n",
        "      \"ds_contrast_delta\": 0.0,\n",
        "      \"ds_saturation_delta\": 0.0,\n",
        "      \"ds_hue_delta\": 0.0,\n",
        "      \"ds_quality_delta\": 0.0,\n",
        "  }\n",
        "  config.models_mutation_ranges = {}\n",
        "  return config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rXGiHWJgnnls"
      },
      "outputs": [],
      "source": [
        "# Configuration for quick tests.\n",
        "def get_config_ti0_cmaterdb():\n",
        "  config = ConfigDict()\n",
        "  config.num_train_examples_between_validations_max = 51200  # 100 batches.\n",
        "  config.num_validations_per_path_training = 2\n",
        "  config.num_validation_examples_max = 5120  # 10 batches.\n",
        "  config.batch_size = 512\n",
        "  config.num_cycles_max = 8\n",
        "  config.num_samples_per_cycle = 2*8\n",
        "  config.max_task_population_size = 7\n",
        "  config.force_mutations = [\"clone:encoder_norm\"]\n",
        "  config.policy_class = \"PPDecay\"\n",
        "  config.policy_kwargs = {}\n",
        "  config.scorer_class = \"ScorerDecay\"\n",
        "  config.scorer_kwargs = dict(\n",
        "      scale_factor=0.99,\n",
        "      num_params=1_000_000,\n",
        "      flops=1_000_000,\n",
        "      )\n",
        "  # The query is used to get the model configs even if the checkpoint is not loaded.\n",
        "  config.vit_checkpoint_query = 'name==\"Ti/16\" and ds==\"i21k\" and aug==\"light1\" and wd==0.1 and sd==0.0'\n",
        "\n",
        "  config.models_default_hparams = {\n",
        "      \"_mu_\": 0.2,\n",
        "      \"num_classes\": 1,\n",
        "      \"mutate_adapters\": True,\n",
        "      \"adapter_layers\": \"\",\n",
        "      \"num_layers\": 0,\n",
        "      \"adapter_dim\": 32,\n",
        "      \"opt_lr\": 0.01,\n",
        "      \"opt_lr_schedule\": \"cosine\",\n",
        "      \"opt_lr_warmup_ratio\": 0.1,\n",
        "      \"opt_momentum\": 0.9,\n",
        "      \"opt_nesterov\": False,\n",
        "      \"ds_image_size\": 32,\n",
        "      \"ds_area_range_min\": 1.0,\n",
        "      \"ds_aspect_ratio_range_min\": 1.0,\n",
        "      \"ds_flip_left_right\": False,\n",
        "      \"ds_brightness_delta\": 0.0,\n",
        "      \"ds_contrast_delta\": 0.0,\n",
        "      \"ds_saturation_delta\": 0.0,\n",
        "      \"ds_hue_delta\": 0.0,\n",
        "      \"ds_quality_delta\": 0.0,\n",
        "  }\n",
        "  config.models_mutation_ranges = {}\n",
        "  return config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "njMI52JMnqq1"
      },
      "outputs": [],
      "source": [
        "def get_config_large():\n",
        "  config = ConfigDict()\n",
        "  config.num_train_examples_between_validations_max = 100_000\n",
        "  config.num_validations_per_path_training = 8\n",
        "  config.num_validation_examples_max = 10_000\n",
        "  config.batch_size = 16\n",
        "  config.num_cycles_max = 24\n",
        "  config.num_samples_per_cycle = 16\n",
        "  config.max_task_population_size = 16\n",
        "  config.force_mutations = [\"clone:encoder_norm\"]\n",
        "  config.policy_class = \"PPDecay\"\n",
        "  config.policy_kwargs = {}\n",
        "  config.scorer_class = \"ScorerDecay\"\n",
        "  config.scorer_kwargs = dict(\n",
        "      scale_factor=0.99,\n",
        "      num_params=2_200_000_000,\n",
        "      flops=3_800_000_000_000,\n",
        "      )\n",
        "  config.vit_checkpoint_query = 'name==\"L/16\" and ds==\"i21k\" and aug==\"medium2\" and wd==0.03 and sd==0.1'\n",
        "\n",
        "  max_num_layers = get_max_num_layers(config.vit_checkpoint_query)\n",
        "  config.models_default_hparams = {\n",
        "      \"_mu_\": 0.2,\n",
        "      \"num_classes\": 1,\n",
        "      \"num_layers\": max_num_layers,\n",
        "      \"opt_lr\": 0.01,\n",
        "      \"opt_lr_schedule\": \"cosine\",\n",
        "      \"opt_lr_warmup_ratio\": 0.05,\n",
        "      \"opt_momentum\": 0.9,\n",
        "      \"opt_nesterov\": False,\n",
        "      \"ds_image_size\": 384,\n",
        "      \"ds_area_range_min\": 1.0,\n",
        "      \"ds_aspect_ratio_range_min\": 1.0,\n",
        "      \"ds_flip_left_right\": False,\n",
        "      \"ds_brightness_delta\": 0.0,\n",
        "      \"ds_contrast_delta\": 0.0,\n",
        "      \"ds_saturation_delta\": 0.0,\n",
        "      \"ds_hue_delta\": 0.0,\n",
        "      \"ds_quality_delta\": 0.0,\n",
        "  }\n",
        "  config.models_mutation_ranges = {}\n",
        "  return config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5_XVkhGhRlp"
      },
      "outputs": [],
      "source": [
        "def config_add_auto_tune(config, agent_class):\n",
        "  \"\"\"Extend config with auto tune parameters.\"\"\"\n",
        "  config.models_mutation_ranges[\"_mu_\"] = [\n",
        "      0.02, 0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18, 0.20, 0.22, 0.24,\n",
        "      0.26, 0.28, 0.30\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"num_layers\"] = list(\n",
        "      range(config.models_default_hparams[\"num_layers\"] + 1))\n",
        "  config.models_mutation_ranges[\"opt_lr\"] = [\n",
        "      0.0001, 0.0002, 0.0005, 0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2,\n",
        "      0.5\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"opt_lr_warmup_ratio\"] = [\n",
        "      0.0, 0.01, 0.02, 0.05, 0.1, 0.2, 0.3\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"opt_momentum\"] = [\n",
        "      0.5, 0.6, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 0.98, 0.99\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"opt_nesterov\"] = [True, False]\n",
        "  config.models_mutation_ranges[\"ds_image_size\"] = [224, 384]\n",
        "  config.models_mutation_ranges[\"ds_area_range_min\"] = [\n",
        "      0.05, 0.5, 0.95, 1.0\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"ds_aspect_ratio_range_min\"] = [\n",
        "      0.5, 0.75, 1.0\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"ds_flip_left_right\"] = [True, False]\n",
        "  config.models_mutation_ranges[\"ds_brightness_delta\"] = [\n",
        "      0.0, 0.01, 0.02, 0.05, 0.1, 0.2\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"ds_contrast_delta\"] = [\n",
        "      0.0, 0.01, 0.02, 0.05, 0.1, 0.2\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"ds_saturation_delta\"] = [\n",
        "      0.0, 0.01, 0.02, 0.05, 0.1, 0.2\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"ds_hue_delta\"] = [\n",
        "      0.0, 0.01, 0.02, 0.05, 0.1, 0.2\n",
        "  ]\n",
        "  config.models_mutation_ranges[\"ds_quality_delta\"] = [\n",
        "      0.0, 0.01, 0.02, 0.05, 0.1, 0.2\n",
        "  ]\n",
        "\n",
        "  # agent_class specific overrides.\n",
        "  if agent_class == \"VitT0\":\n",
        "    config.models_mutation_ranges[\"num_layers\"] = list(range(4))\n",
        "    config.models_mutation_ranges[\"ds_image_size\"] = [32, 48]\n",
        "  elif agent_class in [\"VitT3\", \"VitB\"]:\n",
        "    config.models_mutation_ranges[\"num_layers\"] = list(\n",
        "        range(config.models_default_hparams[\"num_layers\"] + 1))\n",
        "    config.models_mutation_ranges[\"adapter_dim\"] = [8, 16, 32, 64, 128]\n",
        "    config.models_mutation_ranges[\"ds_image_size\"] = [\n",
        "        16 * i for i in (range(1, 1 + int(112 / 16)))\n",
        "    ]\n",
        "  else:\n",
        "    assert agent_class == \"Vit\", (\n",
        "        f\"Undefined mutation ranges for benchmark: {agent_class}\")\n",
        "  return config\n",
        "\n",
        "def config_validate(config):\n",
        "  for khp in config.models_default_hparams:\n",
        "    if khp in config.models_mutation_ranges:\n",
        "      assert config.models_default_hparams[khp] \\\n",
        "          in config.models_mutation_ranges[khp]\n",
        "  for khp in config.models_mutation_ranges:\n",
        "    assert khp in config.models_default_hparams"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "InN16aWQ8F7R"
      },
      "source": [
        "## Agents utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PXtebdPFHhur"
      },
      "outputs": [],
      "source": [
        "def format_agent_id(agent_id: str):\n",
        "  assert \"~\" not in agent_id, f\"Invalid agent id: {agent_id}\"\n",
        "  return agent_id.replace(\"/\", \"~\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5j4LPnTA8EzH"
      },
      "outputs": [],
      "source": [
        "def run_cycles(agent):\n",
        "  config = agent.config\n",
        "  devices = jax.local_devices()\n",
        "  print(\"DEVICE COUNT:\", len(devices))\n",
        "  task_name = config.task_name\n",
        "  num_cycles = config.num_cycles_max\n",
        "\n",
        "  for _ in range(num_cycles):\n",
        "    agent.load_state()\n",
        "    if agent.cycle_id \u003e= num_cycles:\n",
        "      break\n",
        "    print(\"\\n\\n====\")\n",
        "    print(f\"CYCLE: [{agent.cycle_id+1}/{num_cycles}]\")\n",
        "    task = Path.cached_tasks(task_name=task_name)\n",
        "    agent.pop.start_cycle()\n",
        "    agent_cycle(\n",
        "        task, devices, agent.pop, agent.generation_id, agent.cycle_id, config)\n",
        "    agent.pop.end_cycle()\n",
        "    agent.cycle_id += 1\n",
        "    agent.generation_id = 0\n",
        "\n",
        "    write_threads = save_state(\n",
        "        agent.pop, agent.cycle_id, agent.generation_id, config)\n",
        "    # Display stats.\n",
        "    avg_time_per_sample = (\n",
        "        agent.pop.paths_df[\"metrics.end_time\"].mean() \\\n",
        "            - agent.pop.paths_df[\"metrics.start_time_loop\"].mean()\n",
        "        ) / len(devices)\n",
        "    print(f\"Avg time per path: {avg_time_per_sample:.2f} s\")\n",
        "    # Wait for last state write to complete.\n",
        "    for t in write_threads:\n",
        "      t.join()\n",
        "    if agent.cycle_id \u003e= num_cycles:\n",
        "      break\n",
        "\n",
        "def run_root_model(agent):\n",
        "  agent.load_state()\n",
        "  pop = agent.pop\n",
        "  config = agent.config\n",
        "  cycle_id = agent.cycle_id\n",
        "  generation_id = agent.generation_id\n",
        "  write_threads = save_state(pop, cycle_id, generation_id, config)\n",
        "  for t in write_threads:\n",
        "    t.join()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hElrroN38JcE"
      },
      "source": [
        "## Agent classes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_wzXhNikODR8"
      },
      "outputs": [],
      "source": [
        "class Agent():\n",
        "  @property\n",
        "  def class_name(self):\n",
        "    return self.__class__.__name__\n",
        "\n",
        "  @property\n",
        "  def id(self):\n",
        "    return self.config.agent_id\n",
        "\n",
        "  def run(self):\n",
        "    assert False, \"Not implementd\"\n",
        "\n",
        "  def complete_config(self, task_name, experiment_dir):\n",
        "    self.config.task_name = task_name\n",
        "    self.config.experiment_dir = experiment_dir\n",
        "    self.config.agent_id = format_agent_id(f\"{self.class_name}/{task_name}\")\n",
        "    self.config.agent_dir = os.path.join(experiment_dir, self.id)\n",
        "\n",
        "\n",
        "class Vit(Agent):\n",
        "  \"\"\"ViT large / all tasks\"\"\"\n",
        "\n",
        "  def __init__(self, task_name, experiment_dir, auto_tune, scale_factor):\n",
        "    self.config = self.init_config_fn()\n",
        "    self.complete_config(task_name, experiment_dir, auto_tune, scale_factor)\n",
        "\n",
        "  def load_state(self):\n",
        "    task_name = self.config.task_name\n",
        "    self.pop = Population(self.config)\n",
        "    self.cycle_id = 0\n",
        "    self.generation_id = 0\n",
        "    # Root models.\n",
        "    if task_name.startswith(\"root_model/\"):\n",
        "      hparams = self.config.models_default_hparams.as_configdict()\n",
        "      if task_name == \"root_model/random_init\":\n",
        "        path_params = Path.cached_init_params(\n",
        "          query=self.config.vit_checkpoint_query,\n",
        "          **hparams)\n",
        "      else:\n",
        "        assert task_name == \"root_model/checkpoint\", task_name\n",
        "        path_params = get_vit_checkpoint_mapped(\n",
        "            hparams[\"ds_image_size\"],\n",
        "            self.config.vit_checkpoint_query)\n",
        "      path = Path(\n",
        "          hparams,\n",
        "          params2comps(\n",
        "              path_params, train_locks=[self.id], agent_id=self.id),\n",
        "          parent=None,\n",
        "          agent_id=self.id,\n",
        "          task_name=task_name)\n",
        "      self.pop.paths[self.id].append(path)\n",
        "      return\n",
        "\n",
        "    # Load latest agent state.\n",
        "    def validate_df(df):\n",
        "      assert len(df[\"agent_id\"].unique()) == 1, len(df[\"agent_id\"].unique())\n",
        "      assert df[\"agent_id\"].unique()[0] == self.id, df[\"agent_id\"].unique()[0]\n",
        "    agent_checkpoint = latest_checkpoint(\n",
        "        os.path.join(self.config.agent_dir, \"state_*_*/\"))\n",
        "    if agent_checkpoint:\n",
        "      matched = re.findall(r\"checkpoint_([0-9]+)_([0-9]+)$\", agent_checkpoint)\n",
        "      assert len(matched) == 1\n",
        "      self.cycle_id = int(matched[0][0])\n",
        "      self.generation_id = int(matched[0][1])\n",
        "      state_dir = os.path.dirname(agent_checkpoint)\n",
        "      self.pop.paths_df = df_read_from_csv(state_dir, \"paths\")\n",
        "      self.pop.comps_df = df_read_from_csv(state_dir, \"components\")\n",
        "      validate_df(self.pop.paths_df)\n",
        "      validate_df(self.pop.comps_df)\n",
        "      # Set globals.\n",
        "      Path.paths = []\n",
        "      Path.counter = 1 + int(self.pop.paths_df.id.max())\n",
        "      Component.counter = 1 + int(self.pop.comps_df.id.max())\n",
        "      # Get id of the last componet saved in a non intermediate checkpoint.\n",
        "      non_intermediated_checkpoint = latest_checkpoint(\n",
        "          os.path.join(self.config.agent_dir, \"state_*_0/\"))\n",
        "      if non_intermediated_checkpoint:\n",
        "        ni_paths_df = df_read_from_csv(\n",
        "            os.path.dirname(non_intermediated_checkpoint), \"paths\")\n",
        "        validate_df(ni_paths_df)\n",
        "        Path.last_saved = int(ni_paths_df.id.max())\n",
        "        ni_comps_df = df_read_from_csv(\n",
        "            os.path.dirname(non_intermediated_checkpoint), \"components\")\n",
        "        validate_df(ni_comps_df)\n",
        "        Component.last_saved = int(ni_comps_df.id.max())\n",
        "      print(\"CONTINUING FROM STATE\", self.cycle_id, self.generation_id)\n",
        "\n",
        "    # Load all available paths.\n",
        "    all_agents_dirs = os.path.join(self.config.experiment_dir, \"*\")\n",
        "    state_dir = os.path.dirname(agent_checkpoint) if agent_checkpoint else None\n",
        "    load_paths(self.pop, state_dir, all_agents_dirs)\n",
        "\n",
        "    assert self.pop.paths, \"Empty population, run an agent creating a \" \\\n",
        "        \"root model to initialize the population.\"\n",
        "    df_leaderboard(pop_to_df(self.pop))\n",
        "\n",
        "  @property\n",
        "  def init_config_fn(self):\n",
        "    return get_config_large\n",
        "\n",
        "  def run(self):\n",
        "    if self.config.task_name.startswith(\"root_model/\"):\n",
        "      run_root_model(self)\n",
        "      return\n",
        "    run_cycles(self)\n",
        "\n",
        "  def complete_config(self, task_name, experiment_dir, auto_tune, scale_factor):\n",
        "    super().complete_config(task_name, experiment_dir)\n",
        "    if auto_tune:\n",
        "      self.config = config_add_auto_tune(self.config, self.class_name)\n",
        "    if scale_factor:\n",
        "      self.config.scorer_kwargs[\"scale_factor\"] = scale_factor\n",
        "    self.config = FrozenConfigDict(self.config)\n",
        "    config_validate(self.config)\n",
        "\n",
        "\n",
        "class VitT0(Vit):\n",
        "  \"\"\"ViT tiny 0 layers / cmaterdb benchmark\"\"\"\n",
        "  @property\n",
        "  def init_config_fn(self):\n",
        "    return get_config_ti0_cmaterdb\n",
        "\n",
        "\n",
        "class VitT3(Vit):\n",
        "  \"\"\"ViT tiny 3 layers / characters benchmark\"\"\"\n",
        "  @property\n",
        "  def init_config_fn(self):\n",
        "    return get_config_ti3_chars\n",
        "\n",
        "\n",
        "class VitB(Vit):\n",
        "  \"\"\"ViT base / decathlon benchmark\"\"\"\n",
        "  @property\n",
        "  def init_config_fn(self):\n",
        "    return get_config_base_deca"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gfGafBjhDpV7"
      },
      "source": [
        "# ViT Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3MRsK4hvocq8"
      },
      "outputs": [],
      "source": [
        "def get_sample_images(image_size:int, batch_size:int):\n",
        "  return np.zeros((batch_size, image_size, image_size, 3))\n",
        "\n",
        "def get_sample_labels(batch_size:int):\n",
        "  return np.zeros(batch_size, dtype=np.int32)\n",
        "\n",
        "def get_sample_batch(image_size:int, batch_size:int):\n",
        "  return {\"image\": get_sample_images(image_size, batch_size),\n",
        "          \"label\": get_sample_labels(batch_size),}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0jPjzRlYYi0x"
      },
      "outputs": [],
      "source": [
        "def get_vit_checkpoint(image_size, query):\n",
        "  filename = get_vit_filename(query)\n",
        "\n",
        "  config = get_vit_config(query)\n",
        "\n",
        "  model = VisionTransformer(**config, num_classes=2)  # num_classes unsed.\n",
        "  init_params = copy.deepcopy(jax.device_get(\n",
        "      model.init(jax.random.PRNGKey(random.randrange(1e10)),\n",
        "                 get_sample_images(image_size=image_size,\n",
        "                                   batch_size=1),\n",
        "                 train=False  # Disable dropout.\n",
        "                 )[\"params\"]))\n",
        "\n",
        "  params = checkpoint.load_pretrained(\n",
        "    pretrained_path=f\"gs://vit_models/augreg/{filename}.npz\",\n",
        "    init_params=init_params,\n",
        "    model_config=config)\n",
        "\n",
        "  return params\n",
        "\n",
        "def get_vit_checkpoint_mapped(image_size, query):\n",
        "  params = get_vit_checkpoint(image_size, query)\n",
        "  params = params_model_to_comps(params)\n",
        "  return params\n",
        "\n",
        "def get_reshaped_posembed_component(\n",
        "    agent_id: str, ds_image_size: int, query: str):\n",
        "  params = get_vit_checkpoint_mapped(ds_image_size, query)[\"posembed_input\"]\n",
        "  return Component(name=\"posembed_input\",\n",
        "                   agent_id=agent_id,\n",
        "                   params=params,\n",
        "                   train_locks=[])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G_Xsw_tdLIRC"
      },
      "outputs": [],
      "source": [
        "# Parameter mapping.\n",
        "TRANSFORMER_KEYS = set(\n",
        "    [\"encoder_norm\", \"posembed_input\" ] + \\\n",
        "    [f\"encoderblock_{k}\" for k in range(30)])\n",
        "\n",
        "def params_model_to_comps(params):\n",
        "  global TRANSFORMER_KEYS\n",
        "  TRANSFORMER_KEYS.update(params[\"Transformer\"].keys())\n",
        "  new_params = {}\n",
        "  for k in params.keys():\n",
        "    if k == \"Transformer\":\n",
        "      t_params = params[k]\n",
        "      for t_k in t_params.keys():\n",
        "        new_params[t_k] = t_params[t_k]\n",
        "    else:\n",
        "      new_params[k] = params[k]\n",
        "  return flax.core.freeze(new_params)\n",
        "\n",
        "def params_comps_to_model(params):\n",
        "  params = params.unfreeze()\n",
        "\n",
        "  params[\"Transformer\"] = {}\n",
        "  keys = list(params.keys())\n",
        "  assert len(TRANSFORMER_KEYS) != 0\n",
        "  for k in keys:\n",
        "    if k in TRANSFORMER_KEYS:\n",
        "      params[\"Transformer\"][k] = params.pop(k)\n",
        "  return flax.core.freeze(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2Ktk16O9PhYZ"
      },
      "outputs": [],
      "source": [
        "def get_vit_model(\n",
        "    num_classes: int,\n",
        "    num_layers: int,\n",
        "    query: str,\n",
        "    adapter_layers: str = \"\",\n",
        "    adapter_dim: int = -1,\n",
        "    ):\n",
        "  config = get_vit_config(query)\n",
        "  config[\"transformer\"][\"num_layers\"] = num_layers\n",
        "  config[\"transformer\"][\"adapter_layers\"] = adapter_layers\n",
        "  config[\"transformer\"][\"adapter_dim\"] = adapter_dim\n",
        "  config = FrozenConfigDict(config)\n",
        "  model = VisionTransformer(**config, num_classes=num_classes)\n",
        "  return model\n",
        "\n",
        "def get_vit_model_and_params(\n",
        "    num_classes, num_layers, ds_image_size, query,\n",
        "    adapter_layers, adapter_dim):\n",
        "  model = get_vit_model(\n",
        "      num_classes, num_layers, query, adapter_layers, adapter_dim)\n",
        "  init_params = copy.deepcopy(jax.device_get(\n",
        "      model.init(\n",
        "          jax.random.PRNGKey(random.randrange(1e10)),\n",
        "          get_sample_images(image_size=ds_image_size, batch_size=1),\n",
        "          train=False  # Disable dropout.\n",
        "          )[\"params\"]))\n",
        "  return model, init_params\n",
        "\n",
        "def get_vit_params_mapped(\n",
        "    num_classes: int,\n",
        "    num_layers: int,\n",
        "    ds_image_size: int,\n",
        "    query: str,\n",
        "    adapter_layers: str = \"\",\n",
        "    adapter_dim: int = -1):\n",
        "  model, init_params = get_vit_model_and_params(\n",
        "      num_classes, num_layers, ds_image_size, query,\n",
        "      adapter_layers, adapter_dim)\n",
        "  init_params = params_model_to_comps(init_params)\n",
        "  return init_params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-F8s09QiK8ri"
      },
      "outputs": [],
      "source": [
        "def format_params(a, b):\n",
        "  params = a.copy(b)\n",
        "  assert len(params) == len(a) + len(b)  # Dicts keys should not overlap.\n",
        "  params = params_comps_to_model(params)\n",
        "  return params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FsuttUqHvYGr"
      },
      "outputs": [],
      "source": [
        "def get_optimizer(\n",
        "    opt_lr: float,\n",
        "    opt_lr_schedule: str,\n",
        "    opt_lr_warmup_ratio: float,\n",
        "    opt_momentum: float,\n",
        "    opt_nesterov: bool,\n",
        "    num_train_batches_between_validations: int,\n",
        "    num_validations_per_path_training: int,\n",
        "    ):\n",
        "  min_lr = opt_lr / 1000.0\n",
        "  if opt_lr_schedule == \"constant\":\n",
        "    # Divide by 2 so that average lr is the same as other types.\n",
        "    learning_rate = 0.5 * opt_lr\n",
        "  elif opt_lr_schedule == \"linear\":\n",
        "    train_steps = int(num_train_batches_between_validations\n",
        "                      * num_validations_per_path_training)\n",
        "    warmup_steps = int(opt_lr_warmup_ratio * train_steps)\n",
        "    schedules = [\n",
        "        optax.linear_schedule(\n",
        "            init_value=min_lr,\n",
        "            end_value=opt_lr,\n",
        "            transition_steps=warmup_steps),\n",
        "        optax.linear_schedule(\n",
        "            init_value=opt_lr,\n",
        "            end_value=min_lr,\n",
        "            transition_steps=train_steps-warmup_steps)]\n",
        "    learning_rate = optax.join_schedules(schedules, [warmup_steps])\n",
        "  elif opt_lr_schedule == \"cosine\":\n",
        "    train_steps = int(num_train_batches_between_validations\n",
        "                      * num_validations_per_path_training)\n",
        "    learning_rate = optax.warmup_cosine_decay_schedule(\n",
        "        init_value=min_lr,\n",
        "        peak_value=opt_lr,\n",
        "        warmup_steps=int(opt_lr_warmup_ratio * train_steps),\n",
        "        decay_steps=train_steps)\n",
        "  elif opt_lr_schedule == \"restarts\":\n",
        "    train_steps = num_train_batches_between_validations\n",
        "    repeats = num_validations_per_path_training\n",
        "    kwargs = dict(\n",
        "        init_value=min_lr,\n",
        "        peak_value=opt_lr,\n",
        "        warmup_steps=int(opt_lr_warmup_ratio * train_steps),\n",
        "        decay_steps=train_steps,\n",
        "    )\n",
        "    kwargs = [kwargs] * repeats\n",
        "    learning_rate = optax.sgdr_schedule(kwargs)\n",
        "  else:\n",
        "    assert False, f\"Invalid lr schedule: {opt_lr_schedule}\"\n",
        "\n",
        "  return optax.chain(\n",
        "      optax.clip_by_global_norm(1.0),\n",
        "      optax.sgd(\n",
        "          learning_rate=learning_rate,\n",
        "          momentum=opt_momentum,\n",
        "          nesterov=opt_nesterov,\n",
        "          accumulator_dtype=jnp.bfloat16))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kqnhOnk5jbnm"
      },
      "outputs": [],
      "source": [
        "def get_default_splits(tfds_name):\n",
        "  info = get_tfds_builder(tfds_name).info\n",
        "  splits = list(info.splits.keys())\n",
        "  assert \"train\" in splits, splits\n",
        "  splits.remove(\"train\")\n",
        "  used_percent = 0\n",
        "  slice_percent = 5\n",
        "  pp = {}\n",
        "  for k in [\"test\", \"validation\"]:\n",
        "    if k in splits:\n",
        "      pp[k] = k\n",
        "      splits.remove(k)\n",
        "    else:\n",
        "      pp[k] = f\"train[{used_percent}%:{used_percent+slice_percent}%]\"\n",
        "      used_percent += slice_percent\n",
        "  pp[\"train\"] = f\"train[{used_percent}%:]\"\n",
        "  return pp\n",
        "\n",
        "def get_dataset_and_splits(tfds_name: str):\n",
        "  vtab_class = None\n",
        "  if tfds_name in [\"imagenet_v2\", \"cifar10_1\"]:\n",
        "    assert False,  f\"{tfds_name} used as validation set for other tasks.\"\n",
        "\n",
        "  if tfds_name == \"imagenet2012\":\n",
        "    dataset = {\n",
        "        \"train\":\"imagenet2012\", \"validation\":\"imagenet_v2\", \"test\":\"imagenet2012\"}\n",
        "    splits = {\n",
        "        \"train\":\"train\", \"validation\":\"test\", \"test\":\"validation\"}\n",
        "  elif tfds_name == \"cifar100\":\n",
        "    dataset = tfds_name\n",
        "    splits = {\n",
        "        \"train\":\"train[:98%]\", \"validation\":\"train[98%:]\", \"test\":\"test\"}\n",
        "  elif tfds_name == \"cifar10\":\n",
        "    dataset = {\n",
        "        \"train\":\"cifar10\", \"validation\":\"cifar10_1\", \"test\":\"cifar10\"}\n",
        "    splits = {\n",
        "        \"train\":\"train\", \"validation\":\"test\", \"test\":\"test\"}\n",
        "  elif (tfds_name.startswith(\"visual_domain_decathlon/\") or\n",
        "        tfds_name in [\"i_naturalist2017\", \"i_naturalist2018\", \"places365_small\"]):\n",
        "    dataset = tfds_name\n",
        "    # Test has no labels, split validation in half.\n",
        "    splits =  {\n",
        "        \"train\":\"train\", \"validation\":\"validation[:50%]\", \"test\":\"validation[50%:]\"}\n",
        "  elif tfds_name.startswith(\"cmaterdb/\"):\n",
        "    dataset = tfds_name\n",
        "    # Increase size of validation set due to small dataset size.\n",
        "    splits =  {\n",
        "        \"train\":\"train[20%:]\", \"validation\":\"train[:20%]\", \"test\":\"test\"}\n",
        "  elif tfds_name == \"omniglot\":\n",
        "    # Test has no labels, and missing validation, use additional splits.\n",
        "    dataset = tfds_name\n",
        "    splits = {\"train\":\"train\", \"validation\":\"small1\", \"test\":\"small2\"}\n",
        "  elif tfds_name.startswith(\"controlled_noisy_web_labels/\"):\n",
        "    dataset = tfds_name\n",
        "    splits =  {\n",
        "        \"train\":\"train_00\",\n",
        "        \"validation\":\"validation[:50%]\",\n",
        "        \"test\":\"validation[50%:]\"}\n",
        "  elif tfds_name.startswith(\"cycle_gan/\"):\n",
        "    dataset = tfds_name\n",
        "    splits =  {\n",
        "        \"train\":\"trainA[10%:]+trainB[10%:]\",\n",
        "        \"validation\":\"trainA[:10%]+trainB[:10%]\",\n",
        "        \"test\":\"testA+testB\"}\n",
        "  elif tfds_name in [\"imagenet_a\", \"imagenet_r\", \"imagenet_sketch\",\n",
        "                     \"siscore/rotation\", \"siscore/size\", \"siscore/location\",]:\n",
        "    # Only test split.\n",
        "    dataset = tfds_name\n",
        "    splits =  {\n",
        "        \"train\":\"test[10%:]\",\n",
        "        \"validation\":\"test[5%:10%]\",\n",
        "        \"test\":\"test[:5%]\"}\n",
        "  elif tfds_name in [\"pet_finder\"]:\n",
        "    # Explicitly use only train split. E.g. test has no labels.\n",
        "    dataset = tfds_name\n",
        "    splits =  {\n",
        "        \"train\":\"train[10%:]\",\n",
        "        \"validation\":\"train[5%:10%]\",\n",
        "        \"test\":\"train[:5%]\"}\n",
        "  elif tfds_name == \"quickdraw_bitmap\":\n",
        "    dataset = tfds_name\n",
        "    # Cap size of test and validation set.\n",
        "    splits =  {\n",
        "        \"train\":\"train[20000:]\", \"validation\":\"train[10000:20000]\", \"test\":\"train[:10000]\"}\n",
        "  elif tfds_name == \"stanford_online_products\":\n",
        "    dataset = tfds_name\n",
        "    # Use the first 10k test samples as validation since test has 60k.\n",
        "    splits =  {\n",
        "        \"train\":\"train\", \"validation\":\"test[:10000]\", \"test\":\"test[10000:]\"}\n",
        "  elif tfds_name in VTAB_TASKS or (\n",
        "      tfds_name.endswith(\"/1k\") and tfds_name.replace(\"/1k\", \"\") in VTAB_TASKS):\n",
        "    is_vtab_1k = tfds_name.endswith(\"/1k\")\n",
        "    tfds_name = tfds_name.replace(\"/1k\", \"\")\n",
        "    registry_name = {\n",
        "        \"diabetic_retinopathy_detection/btgraham-300\": \"diabetic_retinopathy\",\n",
        "        \"svhn_cropped\": \"svhn\",\n",
        "        \"cifar100\": \"cifar\",\n",
        "        \"cifar10\": \"cifar\",\n",
        "    }.get(tfds_name, tfds_name.split(\"/\")[0])\n",
        "    args = {\n",
        "        \"clevr/count_all\": (\"count_all\",),\n",
        "        \"clevr/count_cylinders\": (\"count_cylinders\",),\n",
        "        \"clevr/closest_object_distance\": (\"closest_object_distance\",),\n",
        "        \"dsprites/label_x_position\": (\"label_x_position\",),\n",
        "        \"dsprites/label_orientation\": (\"label_orientation\",),\n",
        "        \"kitti/closest_object_distance\": (\"closest_object_distance\",),\n",
        "        \"kitti/count_vehicles\": (\"count_vehicles\",),\n",
        "        \"kitti/closest_vehicle_distance\": (\"closest_vehicle_distance\",),\n",
        "        \"smallnorb/label_category\": (\"label_category\",),\n",
        "        \"smallnorb/label_lighting\": (\"label_lighting\",),\n",
        "        \"smallnorb/label_azimuth\": (\"label_azimuth\",),\n",
        "        \"smallnorb/label_elevation\": (\"label_elevation\",),\n",
        "        \"cifar100\": (100,),\n",
        "        \"cifar10\": (10,),\n",
        "    }.get(tfds_name, ())\n",
        "    vtab_class = task_adapt_registry.Registry.lookup(\n",
        "        f\"data.{registry_name}\")(*args)\n",
        "    vtab_splits = vtab_class._tfds_splits\n",
        "    dataset = {\n",
        "        \"caltech101\": \"caltech101:3.*.*\",\n",
        "        \"dtd\": \"dtd:3.*.*\",\n",
        "        \"oxford_flowers102\": \"oxford_flowers102:2.*.*\",\n",
        "        \"oxford_iiit_pet\": \"oxford_iiit_pet:3.*.*\",\n",
        "        \"sun397\": \"sun397/tfds:4.*.*\",\n",
        "        \"svhn\": \"svhn_cropped:3.*.*\",\n",
        "        \"patch_camelyon\": \"patch_camelyon:2.*.*\",\n",
        "        \"eurosat\": \"eurosat/rgb:2.*.*\",\n",
        "        \"resisc45\": \"resisc45:3.*.*\",\n",
        "        \"diabetic_retinopathy\": \"diabetic_retinopathy_detection/btgraham-300:3.*.*\",\n",
        "        \"clevr\": \"clevr:3.*.*\",\n",
        "        \"dmlab\": \"dmlab:2.0.1\",\n",
        "        \"dsprites\": \"dsprites:2.*.*\",\n",
        "        \"kitti\": \"kitti:3.2.0\",\n",
        "        \"smallnorb\": \"smallnorb:2.*.*\",\n",
        "        \"cifar\" : \"cifar100:3.*.*\" if tfds_name == \"cifar100\" else \"cifar10:3.*.*\",\n",
        "    }[registry_name]\n",
        "    if is_vtab_1k:\n",
        "      splits =  {\n",
        "          \"train\": str(vtab_splits[\"train800\"]),\n",
        "          \"validation\": str(vtab_splits[\"val200\"]),\n",
        "          \"test\": str(vtab_splits[\"test\"]),\n",
        "          }\n",
        "    else:\n",
        "      splits =  {\n",
        "          \"train\": str(vtab_splits[\"train\"]),\n",
        "          \"validation\": str(vtab_splits[\"val\"]),\n",
        "          \"test\": str(vtab_splits[\"test\"]),\n",
        "          }\n",
        "  else:\n",
        "    dataset = tfds_name\n",
        "    splits = get_default_splits(tfds_name)\n",
        "  return dataset, splits, vtab_class\n",
        "\n",
        "\n",
        "class Task():\n",
        "  def __init__(self, name, config):\n",
        "    self.config = config\n",
        "\n",
        "    self.dataset, self.splits, self.vtab_class = get_dataset_and_splits(name)\n",
        "    self.name = name\n",
        "    if self.vtab_class:\n",
        "      self.num_classes = self.vtab_class.get_num_classes()\n",
        "    else:\n",
        "      self.num_classes = self.get_builder(\n",
        "          \"train\").info.features[self.get_label_key()].num_classes\n",
        "    num_train_examples = self.get_builder(\n",
        "        \"train\").info.splits[self.splits[\"train\"]].num_examples\n",
        "    self.train_batch_size = config.batch_size\n",
        "    self.num_train_batches_between_validations = math.ceil(\n",
        "        min(num_train_examples,\n",
        "            config.num_train_examples_between_validations_max)\n",
        "        / self.train_batch_size)\n",
        "    self.cache_train = num_train_examples \u003c min(100_000, (\n",
        "        config.num_validations_per_path_training\n",
        "        * self.num_train_batches_between_validations\n",
        "        * self.train_batch_size))\n",
        "\n",
        "    num_validation_examples_tot = self.get_builder(\n",
        "        \"validation\").info.splits[self.splits[\"validation\"]].num_examples\n",
        "    if config.num_validation_examples_max \u003c= num_validation_examples_tot:\n",
        "      self.validation_batch_size = config.batch_size\n",
        "      self.num_validation_batches = math.floor(\n",
        "          config.num_validation_examples_max / self.validation_batch_size)\n",
        "    else:\n",
        "      # Adjust batch_size and num_batches to cover the smaller validation sets.\n",
        "      self.num_validation_batches = math.ceil(\n",
        "          num_validation_examples_tot / config.batch_size)\n",
        "      self.validation_batch_size = math.floor(\n",
        "          num_validation_examples_tot / self.num_validation_batches)\n",
        "      assert num_validation_examples_tot \u003e= (\n",
        "          self.num_validation_batches*self.validation_batch_size)\n",
        "    self.num_validation_examples = (\n",
        "        self.num_validation_batches * self.validation_batch_size)\n",
        "\n",
        "    print(f\"Task: {self.name}\")\n",
        "    print(f\"  Train batches between validations: {self.num_train_batches_between_validations}\")\n",
        "    print(f\"  Validation batches: {self.num_validation_batches}\")\n",
        "    print(f\"  Validation batch size: {self.validation_batch_size}\")\n",
        "    print(f\"  Dataset {{\\n{self.dataset}}}\")\n",
        "    print(f\"  Splits {{\\n{self.splits}}}\")\n",
        "\n",
        "  def get_label_key(self):\n",
        "    return {\n",
        "        \"stanford_online_products\": \"super_class_id\",\n",
        "        }.get(self.name, \"label\")\n",
        "\n",
        "  def get_builder(self, mode):\n",
        "    if type(self.dataset) == str:\n",
        "      return get_tfds_builder(self.dataset)\n",
        "    return get_tfds_builder(self.dataset[mode])\n",
        "\n",
        "  def __str__(self):\n",
        "    return f\"Task_{self.name}\"\n",
        "\n",
        "  def get_ds(self, mode, hparams):\n",
        "    data = self.get_builder(mode).as_dataset(\n",
        "        split=self.splits[mode],\n",
        "        shuffle_files=mode=='train')\n",
        "\n",
        "    def _pp(data):\n",
        "      im = data[\"image\"]\n",
        "      tf.debugging.assert_type(im, tf.uint8)\n",
        "\n",
        "      if mode == \"train\":\n",
        "        if hparams.get(\"ds_quality_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_jpeg_quality(\n",
        "              im,\n",
        "              min_jpeg_quality=int(100 * (1 - hparams[\"ds_quality_delta\"])),\n",
        "              max_jpeg_quality=100)\n",
        "\n",
        "      # Must have 3 channels.\n",
        "      if im.shape[-1] == 1:\n",
        "        im = tf.squeeze(tf.stack([im] * 3, -1), axis=-2)\n",
        "      assert im.shape[-1] == 3\n",
        "\n",
        "      if mode == \"train\":\n",
        "        if hparams.get(\"ds_area_range_min\", 1.0) \u003c 1.0:\n",
        "          channels = im.shape[-1]\n",
        "          begin, size, _ = tf.image.sample_distorted_bounding_box(\n",
        "              tf.shape(im),\n",
        "              tf.zeros([0, 0, 4], tf.float32),\n",
        "              aspect_ratio_range=[hparams[\"ds_aspect_ratio_range_min\"],\n",
        "                                  1.0/hparams[\"ds_aspect_ratio_range_min\"]],\n",
        "              area_range=[hparams[\"ds_area_range_min\"], 1.0],\n",
        "              # Overlap with bounding box, the bounding box should anyway\n",
        "              # default defaults to whole image in this case.\n",
        "              min_object_covered=0,\n",
        "              use_image_if_no_bounding_boxes=True)\n",
        "          im = tf.slice(im, begin, size)\n",
        "          # Restore the depth-dimension lost by the above operation.\n",
        "          im.set_shape([None, None, channels])\n",
        "        if hparams.get(\"ds_flip_left_right\", False):\n",
        "          if tf.random.uniform(shape=[]) \u003e 0.5:\n",
        "            im = tf.image.flip_left_right(im)\n",
        "        if hparams.get(\"ds_brightness_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_brightness(\n",
        "              im, max_delta=hparams[\"ds_brightness_delta\"])\n",
        "        if hparams.get(\"ds_contrast_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_contrast(\n",
        "              im, lower=1 - hparams[\"ds_contrast_delta\"],\n",
        "              upper=1 + hparams[\"ds_contrast_delta\"])\n",
        "        if hparams.get(\"ds_saturation_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_saturation(\n",
        "              im, lower=1 - hparams[\"ds_saturation_delta\"],\n",
        "              upper=1 + hparams[\"ds_saturation_delta\"])\n",
        "        if hparams.get(\"ds_hue_delta\", 0.0) \u003e 0.0:\n",
        "          im = tf.image.random_hue(im, max_delta=hparams[\"ds_hue_delta\"])\n",
        "\n",
        "      tf.debugging.assert_type(im, tf.uint8)\n",
        "      im = tf.image.resize(im, [hparams[\"ds_image_size\"],\n",
        "                                hparams[\"ds_image_size\"]])\n",
        "      tf.debugging.assert_type(im, tf.float32)\n",
        "\n",
        "      # Values in range [-1 , 1].\n",
        "      im = im / 127.5 - 1\n",
        "      im = tf.clip_by_value(im, -1, 1)\n",
        "\n",
        "      return {\"image\": im, \"label\": data[self.get_label_key()]}\n",
        "\n",
        "    if mode == \"validation\":\n",
        "      data = data.take(self.num_validation_examples)\n",
        "    if mode == \"validation\" or (mode == \"train\" and self.cache_train):\n",
        "      data = data.cache()\n",
        "    if mode != \"test\":\n",
        "      data = data.repeat()\n",
        "    if self.vtab_class and self.vtab_class._base_preprocess_fn:\n",
        "      data = data.map(self.vtab_class._base_preprocess_fn, tf.data.AUTOTUNE)\n",
        "    data = data.map(_pp, tf.data.AUTOTUNE)\n",
        "    if mode == \"train\":\n",
        "      batch_size = self.train_batch_size\n",
        "    else:\n",
        "      batch_size = self.validation_batch_size\n",
        "    data = data.batch(batch_size)\n",
        "    if mode == \"train\":\n",
        "      data = data.shuffle(10)\n",
        "    return tfds.as_numpy(data.prefetch(tf.data.AUTOTUNE))\n",
        "\n",
        "def get_task_factory_fn(config):\n",
        "  def get_task(task_name: str):\n",
        "    return Task(name=task_name, config=config)\n",
        "  return get_task"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2A3aZXSWYwCd"
      },
      "outputs": [],
      "source": [
        "def get_num_params(params):\n",
        "  return sum(jax.tree_util.tree_flatten(\n",
        "      jax.tree_util.tree_map(lambda p: np.prod(p.shape), params)\n",
        "      )[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rzxoZ4rdQZcA"
      },
      "outputs": [],
      "source": [
        "def params2comps(params, train_locks, agent_id, name=None):\n",
        "  \"\"\"Convert frozend dict of params to a list of components.\"\"\"\n",
        "  components = []\n",
        "  for k in params:\n",
        "    if name is None or name == k:\n",
        "      c = Component(\n",
        "          name=k, agent_id=agent_id,\n",
        "          params=params[k], train_locks=train_locks)\n",
        "      components.append(c)\n",
        "  return components\n",
        "\n",
        "def params2comp_names(params):\n",
        "  return list(params.keys())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DNxjDX13_dm_"
      },
      "outputs": [],
      "source": [
        "def fingerprint_params(params):\n",
        "  return np.sum(np.array(jax.tree_util.tree_leaves(\n",
        "      jax.tree_util.tree_map(jnp.sum, params))))\n",
        "\n",
        "class Component():\n",
        "  counter = 0\n",
        "  # Components of retained paths with id \u003c= last_saved are saved in checkpoint.\n",
        "  last_saved = -1\n",
        "\n",
        "  def reset_globals():\n",
        "    Component.counter = 0\n",
        "    Component.last_saved = -1\n",
        "\n",
        "  def __init__(\n",
        "      self, name: str, agent_id: str, params, train_locks, opt_state=None):\n",
        "    self.name = name\n",
        "    self.agent_id = agent_id\n",
        "    self.params = jax.device_get(params)\n",
        "    self.opt_state = jax.device_get(opt_state)\n",
        "    self.num_params = None\n",
        "    self.train_locks = set(train_locks)\n",
        "    self.id = Component.counter\n",
        "    Component.counter += 1\n",
        "\n",
        "  def __str__(self):\n",
        "    rtn = f\"Component: {self.id}\\n  Name: {self.name}\"\n",
        "    rtn += f\"\\n  Train locks: {self.train_locks}\"\n",
        "    rtn += f\"\\n  Fingerprint: {self.fingerprint()}\"\n",
        "    rtn += f\"\\n  Num params: {self.get_num_params()}\"\n",
        "    return rtn\n",
        "\n",
        "  def get_num_params(self):\n",
        "    if self.num_params is None:\n",
        "      self.num_params = get_num_params(self.params)\n",
        "    return self.num_params\n",
        "\n",
        "  def fingerprint(self):\n",
        "    return fingerprint_params(self.params)\n",
        "\n",
        "  def is_trainable(self):\n",
        "    return len(self.train_locks) == 0\n",
        "\n",
        "  def clone(self, agent_id):\n",
        "    return Component(name=self.name,\n",
        "                     agent_id=agent_id,\n",
        "                     params=copy.deepcopy(jax.device_get(self.params)),\n",
        "                     train_locks=set(),\n",
        "                     opt_state=copy.deepcopy(self.opt_state))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2FrDHFPU6NV-"
      },
      "outputs": [],
      "source": [
        "class ObjectCache():\n",
        "  def __init__(self, factory_fn, max_size=None):\n",
        "    self.factory_fn = factory_fn\n",
        "    self.factory_fn_signature = inspect.signature(factory_fn)\n",
        "    self.cache = {}\n",
        "    self.max_size = max_size\n",
        "\n",
        "  def __call__(self, *args, **kwargs):\n",
        "    assert not args, \"No positional arguments allowed.\"\n",
        "    kw_params = {}\n",
        "    fn_name = self.factory_fn.__name__\n",
        "    fn_params = inspect.signature(self.factory_fn).parameters\n",
        "    for k_param, v_param in fn_params.items():\n",
        "      if k_param in kwargs:\n",
        "        kw_params[k_param] = kwargs[k_param]\n",
        "      elif v_param.default != v_param.empty:\n",
        "        # Fallback to declared defalut value.\n",
        "        kw_params[k_param] = fn_params[k_param].default\n",
        "      else:\n",
        "        assert False, (\n",
        "            f\"Missing value for argument {k_param} for function {fn_name}\")\n",
        "\n",
        "      if v_param.annotation != v_param.empty:\n",
        "        # Apply annotated type.\n",
        "        assert isinstance(type(v_param.annotation), type)\n",
        "        kw_params[k_param] = v_param.annotation(kw_params[k_param])\n",
        "\n",
        "    key = json.dumps(kw_params, sort_keys=True)\n",
        "    if key not in self.cache:\n",
        "      if self.max_size and self.max_size \u003c= len(self.cache):\n",
        "        rm_key = random.choice(list(self.cache.keys()))\n",
        "        print(f\"Removed from cache: {fn_name}({rm_key})  [cache size {len(self.cache)}]\")\n",
        "        rm_obj = self.cache.pop(rm_key)\n",
        "        del rm_obj\n",
        "      self.cache[key] = self.factory_fn(**kw_params)\n",
        "      print(f\"Added to cache: {fn_name}({key})  [cache size {len(self.cache)}]\")\n",
        "    return self.cache[key]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kwhVlaqRHp-0"
      },
      "outputs": [],
      "source": [
        "def incremental_mutation(value, values_list:list):\n",
        "  assert value in values_list, f\"{value} not in {values_list}\"\n",
        "  idx = values_list.index(value)\n",
        "  idx += 1 if np.random.uniform() \u003c 0.5 else -1\n",
        "  idx = max(0, min(len(values_list)-1, idx))\n",
        "  return values_list[idx]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BdErPqPuqW9s"
      },
      "outputs": [],
      "source": [
        "def compute_flops_hlo(flax_module, *a, **kw):\n",
        "  \"\"\"Compute FLOPs in flax_module.\"\"\"\n",
        "  # Compute flops on cpu for cross platform consistency.\n",
        "  analysis = jax.jit(flax_module, backend='cpu').lower(*a, **kw).cost_analysis()\n",
        "  return analysis[\"flops\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8PU5ffvd_gC9"
      },
      "outputs": [],
      "source": [
        "class Path():\n",
        "  def reset_globals(config):\n",
        "    Path.config = config\n",
        "    Path.counter = 0\n",
        "    Path.last_saved = -1\n",
        "    Path.paths = []\n",
        "    Path.scorer = globals()[config.scorer_class](**config.scorer_kwargs)\n",
        "    # Cache output of functions calls with same args.\n",
        "    Path.cached_tasks = ObjectCache(get_task_factory_fn(config))\n",
        "    Path.cached_posembed_components = ObjectCache(get_reshaped_posembed_component)\n",
        "    Path.cached_optimizers = ObjectCache(get_optimizer)\n",
        "    Path.cached_models = ObjectCache(get_vit_model)\n",
        "    Path.cached_init_params = ObjectCache(get_vit_params_mapped, max_size=10)\n",
        "\n",
        "  def __init__(self, hparams, components, parent, agent_id, task_name):\n",
        "    self.components = components\n",
        "    self.id = Path.counter\n",
        "    Path.counter += 1\n",
        "    self.agent_id = agent_id\n",
        "    self.task_name = task_name\n",
        "    self.parent = parent\n",
        "    self.hparams = hparams\n",
        "    self.metrics = {\n",
        "        \"generation\": 0 if parent is None else parent.metrics[\"generation\"] + 1,\n",
        "    }\n",
        "    Path.paths.append(self)\n",
        "\n",
        "  def __str__(self):\n",
        "    rtn = f\"Path: {self.id}\"\n",
        "    rtn += f\"\\n  Components: {[c.id for c in self.components]}\"\n",
        "    if self.parent:\n",
        "      rtn += f\"\\n  Parent: {self.parent.id}\"\n",
        "    rtn += f\"\\n  Agent: {self.agent_id}\"\n",
        "    rtn += f\"\\n  Task: {self.task_name}\"\n",
        "    for k,v in self.hparams.items():\n",
        "      rtn += f\"\\n    {k}: {v}\"\n",
        "    for k,v in self.metrics.items():\n",
        "      rtn += f\"\\n    {k}: {v}\"\n",
        "    rtn += f\"\\n  Score: {self.score()}\"\n",
        "    return rtn\n",
        "\n",
        "  @property\n",
        "  def task(self):\n",
        "    return Path.cached_tasks(task_name=self.task_name)\n",
        "\n",
        "  @property\n",
        "  def model(self):\n",
        "    return Path.cached_models(query=self.config.vit_checkpoint_query,\n",
        "                              **self.hparams)\n",
        "\n",
        "  def score(self):\n",
        "    return Path.scorer.score(self)\n",
        "\n",
        "  def get_all_params(self):\n",
        "    params = {}\n",
        "    for c in self.components:\n",
        "      params[c.name] = c.params\n",
        "    return flax.core.freeze(params)\n",
        "\n",
        "  def get_trainable_params(self):\n",
        "    params = {}\n",
        "    for c in self.components:\n",
        "      if c.is_trainable():\n",
        "        params[c.name] = c.params\n",
        "    return flax.core.freeze(params)\n",
        "\n",
        "  def get_fixed_params(self):\n",
        "    params = {}\n",
        "    for c in self.components:\n",
        "      if not c.is_trainable():\n",
        "        params[c.name] = c.params\n",
        "    return flax.core.freeze(params)\n",
        "\n",
        "  def update_trainable(self, trained_params, opt_state):\n",
        "    assert len(trained_params.keys()) == len(opt_state.keys())\n",
        "    trainable_count = 0\n",
        "    for c in self.components:\n",
        "      if c.is_trainable():\n",
        "        trainable_count += 1\n",
        "        assert c.name in trained_params.keys()\n",
        "        assert c.name in opt_state.keys()\n",
        "        c.params = trained_params[c.name]\n",
        "        c.opt_state = opt_state[c.name]\n",
        "    assert len(trained_params.keys()) == trainable_count, (\n",
        "        f\"{len(trained_params.keys())} {trainable_count}\")\n",
        "\n",
        "  def get_num_accounted_params(self):\n",
        "    rtn = 0\n",
        "    for c in self.components:\n",
        "      tl = copy.copy(c.train_locks)\n",
        "      assert type(tl) is set\n",
        "      tl.add(self.agent_id)\n",
        "      assert tl\n",
        "      rtn += c.get_num_params() / len(tl)\n",
        "    return rtn\n",
        "\n",
        "  def get_flops(self):\n",
        "    return compute_flops_hlo(\n",
        "          self.model.apply,\n",
        "          {\"params\": format_params(\n",
        "              self.get_trainable_params(),\n",
        "              self.get_fixed_params(),\n",
        "            )},\n",
        "          get_sample_images(\n",
        "              self.hparams[\"ds_image_size\"],\n",
        "              batch_size=1,\n",
        "          ),\n",
        "          train=False)\n",
        "\n",
        "  def clone(self, ds_hparams, policy):\n",
        "    config = Path.config\n",
        "    agent_id = config.agent_id\n",
        "    task_name = config.task_name\n",
        "    comps = []\n",
        "    new_hparams = copy.deepcopy(self.hparams)\n",
        "    # Overwrite dataset hparams with those sampled for the generation batch.\n",
        "    new_hparams.update(ds_hparams)\n",
        "\n",
        "    def get_component_ref(c, clone):\n",
        "      if c.is_trainable() or clone:\n",
        "        # Clone trainable component.\n",
        "        return c.clone(agent_id=agent_id)\n",
        "      # Refer to frozen component.\n",
        "      return c\n",
        "\n",
        "    for k in sorted(config.models_mutation_ranges):\n",
        "      if ((k in [\"_mu_\", \"num_layers\", \"adapter_dim\"]\n",
        "            or k.startswith(OPTIMIZER_HPARAMS_KEYS_PRERFIX)) and\n",
        "          policy.do_mutate(new_hparams, f\"hp:{k}\")):\n",
        "        new_hparams[k] = incremental_mutation(\n",
        "            new_hparams[k],\n",
        "            config.models_mutation_ranges[k])\n",
        "    if \"adapter_layers\" in new_hparams or new_hparams.get(\"mutate_adapters\", False):\n",
        "      new_hparams[\"adapter_layers\"] = mutate_adapters(\n",
        "          hparams=new_hparams,\n",
        "          policy=policy)\n",
        "\n",
        "    init_params = Path.cached_init_params(\n",
        "        query = config.vit_checkpoint_query,\n",
        "        **new_hparams)\n",
        "    new_comp_names = params2comp_names(init_params)\n",
        "    for new_comp_name in new_comp_names:\n",
        "      comp = None\n",
        "      # Attept to reuse matching componenent from closer ancestor.\n",
        "      ancestor = self\n",
        "      while ancestor is not None:\n",
        "        comps_lookup = {c.name:c for c in ancestor.components}\n",
        "        if new_comp_name in comps_lookup:\n",
        "          # Head must be trainable if no acestor is of same task will fall back\n",
        "          # to random init of correct shape.\n",
        "          if new_comp_name == \"head\" and not comps_lookup[new_comp_name].is_trainable():\n",
        "            assert agent_id != ancestor.agent_id, f\"{agent_id} != {ancestor.agent_id}\"\n",
        "            ancestor = ancestor.parent\n",
        "            continue\n",
        "\n",
        "          # Check shapes match otherwise skip.\n",
        "          if (jax.tree_util.tree_map(jnp.shape, init_params[new_comp_name]) !=\n",
        "              jax.tree_util.tree_map(jnp.shape, comps_lookup[new_comp_name].params)):\n",
        "            if new_comp_name == \"posembed_input\":\n",
        "              # Change of image size changed shape of position embeddings,\n",
        "              # this can happend if ds_image_size is tuned,\n",
        "              # continue searching through ancestors for matching size.\n",
        "              assert \"ds_image_size\" in config.models_mutation_ranges\n",
        "              assert new_hparams[\"ds_image_size\"] != ancestor.hparams[\"ds_image_size\"]\n",
        "              ancestor = ancestor.parent\n",
        "              continue\n",
        "            if new_comp_name.startswith(\"residual_adapter_\"):\n",
        "              # Change of adapter inner dimension changed shape of dense layers,\n",
        "              # this can happend if adapter_dim is tuned,\n",
        "              # continue searching through ancestors for matching size.\n",
        "              assert \"adapter_dim\" in config.models_mutation_ranges\n",
        "              assert new_hparams[\"adapter_dim\"] != ancestor.hparams[\"adapter_dim\"]\n",
        "              ancestor = ancestor.parent\n",
        "              continue\n",
        "\n",
        "            print(f\"WARNING: Shapes do not match for component: {new_comp_name}  {ancestor.agent_id}-\u003e{agent_id}\")\n",
        "            print(jax.tree_util.tree_map(jnp.shape, init_params[new_comp_name]))\n",
        "            print(jax.tree_util.tree_map(jnp.shape, comps_lookup[new_comp_name].params))\n",
        "            assert False  # Should not happen in current configuration.\n",
        "\n",
        "          ancestor_comp = comps_lookup[new_comp_name]\n",
        "          comp = get_component_ref(\n",
        "              ancestor_comp, clone=(\n",
        "                  ancestor_comp.is_trainable() or policy.do_mutate(\n",
        "                      new_hparams, f\"clone:{new_comp_name}\")))\n",
        "          break\n",
        "        ancestor = ancestor.parent\n",
        "\n",
        "      # Get reshaped posembed_input.\n",
        "      if comp is None and new_comp_name == \"posembed_input\":\n",
        "        pe_comp = Path.cached_posembed_components(\n",
        "            agent_id=agent_id,\n",
        "            query=config.vit_checkpoint_query,\n",
        "            **new_hparams)\n",
        "        # Clone to make the component trainable.\n",
        "        comp = get_component_ref(pe_comp, clone=True)\n",
        "\n",
        "      # Otherwise create one from random init params.\n",
        "      if comp is None:\n",
        "        if VERBOSE:\n",
        "          print(\"Init:\", new_comp_name)\n",
        "        # Possible rand init triggering combinations in current configurations.\n",
        "        assert (\n",
        "            new_comp_name == \"head\"\n",
        "            or new_comp_name.startswith(\"residual_adapter_\")\n",
        "            or (new_comp_name.startswith(\"encoderblock_\") and \\\n",
        "                config.models_default_hparams[\"num_layers\"] \u003c max(\n",
        "                config.models_mutation_ranges.get(\"num_layers\", [-1])))\n",
        "            )\n",
        "        comp = params2comps(\n",
        "            init_params, train_locks=[],\n",
        "            agent_id=agent_id, name=new_comp_name)[0]\n",
        "      assert comp is not None\n",
        "      comps.append(comp)\n",
        "\n",
        "    rtn = Path(\n",
        "        new_hparams, comps, parent=self, agent_id=agent_id, task_name=task_name)\n",
        "\n",
        "    if agent_id == self.agent_id:\n",
        "      self.metrics[\"offsprings\"] = self.metrics.get(\"offsprings\", 0) + 1\n",
        "\n",
        "    return rtn\n",
        "\n",
        "  def get_optimizer(self):\n",
        "    return Path.cached_optimizers(\n",
        "        num_train_batches_between_validations=\n",
        "            self.task.num_train_batches_between_validations,\n",
        "        num_validations_per_path_training=\n",
        "            self.task.config.num_validations_per_path_training,\n",
        "        **self.hparams)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C3-3N-LPV09d"
      },
      "outputs": [],
      "source": [
        "def mutate_adapters(hparams, policy, allow_removal=False):\n",
        "  num_layers = hparams[\"num_layers\"]\n",
        "  a_ids = set(ids_str2ints(hparams.get(\"adapter_layers\", \"\")))\n",
        "  if hparams.get(\"mutate_adapters\", False):\n",
        "    for a_id in range(num_layers):\n",
        "      if a_id in a_ids:\n",
        "        if allow_removal and policy.do_mutate(\n",
        "            hparams, f\"remove:residual_adapter_{a_id}\"):\n",
        "          a_ids.remove(a_id)\n",
        "      else:\n",
        "        if policy.do_mutate(hparams, f\"add:residual_adapter_{a_id}\"):\n",
        "          a_ids.add(a_id)\n",
        "  # Drop adapters of layers dropped by a possible mutation in num_layers.\n",
        "  a_ids = [a_id for a_id in a_ids if a_id \u003c num_layers]\n",
        "  return ids_ints2str(a_ids)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bkRmcJgzUbwN"
      },
      "outputs": [],
      "source": [
        "class ScorerDecay():\n",
        "  def __init__(self, scale_factor, num_params, flops=0):\n",
        "    assert 0.0 \u003c scale_factor \u003c= 1.0\n",
        "    self.scale_factor = scale_factor\n",
        "    self.num_params = num_params\n",
        "    self.flops = flops\n",
        "\n",
        "  def score(self, path):\n",
        "    if (\"quality\" not in path.metrics\n",
        "        or math.isnan(path.metrics[\"quality\"])):\n",
        "      return None\n",
        "    assert path.metrics[\"quality\"] \u003e= 0, (\n",
        "        f\"{path.task_name} {path.metrics['quality']}\")\n",
        "    score = path.metrics[\"quality\"]\n",
        "    if self.num_params \u003e 0:\n",
        "      # Accounted params needs to be updated since it depends on the\n",
        "      # changing structure of the system.\n",
        "      path.metrics[\"accounted_params\"] = path.get_num_accounted_params()\n",
        "      score *= self.scale_factor ** (\n",
        "          path.metrics[\"accounted_params\"] / self.num_params)\n",
        "    if self.flops \u003e 0:\n",
        "      if \"flops\" not in path.metrics:\n",
        "        path.metrics[\"flops\"] = path.get_flops()\n",
        "      score *= self.scale_factor ** (path.metrics[\"flops\"] / self.flops)\n",
        "    assert score \u003e= 0\n",
        "    path.metrics[\"score\"] = score\n",
        "    return score"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iWZeLq5zZlfb"
      },
      "outputs": [],
      "source": [
        "class PPDecay():\n",
        "  def __init__(self, config):\n",
        "    self.config = config\n",
        "\n",
        "  def do_mutate(self, hparams, mutation_name):\n",
        "    \"\"\"Returns True if mutation is sampled to be applied.\"\"\"\n",
        "    if mutation_name in self.config.get(\"force_mutations\", []):\n",
        "      return True\n",
        "    mutation_prob_k = f\"_mu_|{mutation_name}\"\n",
        "    # Fallback is used for batch shared sampling.\n",
        "    mu = hparams.get(\"_mu_\", self.config.models_default_hparams[\"_mu_\"])\n",
        "    mutation_prob = hparams.get(mutation_prob_k, mu)\n",
        "    if \"_mu_\" in self.config.models_mutation_ranges:\n",
        "      if mu \u003e np.random.uniform():\n",
        "        mutation_prob = incremental_mutation(\n",
        "            mutation_prob, self.config.models_mutation_ranges[\"_mu_\"])\n",
        "      hparams[mutation_prob_k] = mutation_prob\n",
        "    return mutation_prob \u003e np.random.uniform()\n",
        "\n",
        "  def sample_path(self, pop, ds_hparams):\n",
        "    parent = None\n",
        "    for path in sorted(pop.paths[self.config.agent_id],\n",
        "                       key=lambda p: p.score(),\n",
        "                       reverse=True):\n",
        "      offsprings = path.metrics.get(\"offsprings\", 0)\n",
        "      print(\" \", path.id, int(offsprings))\n",
        "      assert not math.isnan(offsprings)\n",
        "      if np.random.uniform() \u003c 0.5 ** offsprings:\n",
        "        parent = path\n",
        "        break\n",
        "\n",
        "    if not parent:  # Random sample.\n",
        "      parent = random.choice([p for paths in pop.paths.values() for p in paths])\n",
        "      print(\"  random\", parent.id, parent.agent_id)\n",
        "\n",
        "    child = parent.clone(ds_hparams, self)\n",
        "\n",
        "    gc.collect()\n",
        "\n",
        "    # Store record of mutations.\n",
        "    mutations = {}\n",
        "    for k in child.hparams:\n",
        "      if parent.hparams.get(k) != child.hparams[k]:\n",
        "        mutations[k] = (parent.hparams.get(k), child.hparams[k])\n",
        "    child.metrics[\"mutations\"] = json.dumps(mutations)\n",
        "    print(child.id, child.metrics[\"mutations\"])\n",
        "    return child\n",
        "\n",
        "  def sample_ds_hparams(self, pop):\n",
        "    \"\"\"Sample hparams that need to be shared across each paths generation.\"\"\"\n",
        "    assert pop.config is self.config\n",
        "    ds_hparams = {}\n",
        "\n",
        "    # Initialize shared hparams with defaults.\n",
        "    for key in self.config.models_default_hparams:\n",
        "      if key.startswith(DATASET_HPARAMS_KEYS_PRERFIX):\n",
        "        ds_hparams[key] = self.config.models_default_hparams[key]\n",
        "\n",
        "    # Overwrite with values from best path if available.\n",
        "    best_path = pop.get_best_path()\n",
        "    if best_path:\n",
        "      ds_hparams.update(\n",
        "          {k : best_path.hparams[k] for k in ds_hparams if k in best_path.hparams})\n",
        "      ds_hparams.update(\n",
        "          {k : best_path.hparams[k] for k in best_path.hparams if k.startswith(\n",
        "              f\"_mu_|hp:{DATASET_HPARAMS_KEYS_PRERFIX}\")})\n",
        "\n",
        "      # Sample mutations.\n",
        "      for k in list(ds_hparams):\n",
        "        if (k in self.config.models_mutation_ranges\n",
        "            and pop.policy.do_mutate(ds_hparams, f\"hp:{k}\")):\n",
        "          ds_hparams[k] = incremental_mutation(\n",
        "              ds_hparams[k],\n",
        "              self.config.models_mutation_ranges[k])\n",
        "\n",
        "    for k in ds_hparams:\n",
        "      assert (k.startswith(DATASET_HPARAMS_KEYS_PRERFIX) or\n",
        "              k.startswith(f\"_mu_|hp:{DATASET_HPARAMS_KEYS_PRERFIX}\"))\n",
        "    return ds_hparams"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YMbYgKd8_nyi"
      },
      "outputs": [],
      "source": [
        "class Population():\n",
        "  def __init__(self, config):\n",
        "    Path.reset_globals(config)\n",
        "    Component.reset_globals()\n",
        "    self.paths = defaultdict(list)\n",
        "    self.config = config\n",
        "    self.paths_df = pd.DataFrame()\n",
        "    self.comps_df = pd.DataFrame()\n",
        "    self.policy = globals()[config.policy_class](\n",
        "        **config.policy_kwargs,\n",
        "        config=config)\n",
        "\n",
        "  def get_best_path(self):\n",
        "    if len(self.paths[self.config.agent_id]) == 0:\n",
        "      return None\n",
        "    # Most recent path achieving max score.\n",
        "    return max(sorted(\n",
        "        self.paths[self.config.agent_id], key=lambda p: p.id, reverse=True),\n",
        "        key=lambda p: p.score())\n",
        "\n",
        "  def prune_population(self):\n",
        "    if self.config.get(\"max_task_population_size\", None) and (\n",
        "        len(self.paths[self.config.agent_id]) \u003e self.config.max_task_population_size):\n",
        "      best_path = self.get_best_path()\n",
        "      self.paths[self.config.agent_id] = sorted(\n",
        "          self.paths[self.config.agent_id], key=lambda p: p.score(), reverse=True\n",
        "          )[:self.config.max_task_population_size]\n",
        "      assert best_path in self.paths[self.config.agent_id] or best_path.score() == self.get_best_path().score()\n",
        "\n",
        "  def sample_path(self, ds_hparams):\n",
        "    return self.policy.sample_path(pop=self, ds_hparams=ds_hparams)\n",
        "\n",
        "  def sample_ds_hparams(self):\n",
        "    return self.policy.sample_ds_hparams(pop=self)\n",
        "\n",
        "  def add_train_locks(self):\n",
        "    # Check.\n",
        "    for ps in self.paths.values():\n",
        "      for p in ps:\n",
        "        for c in p.components:\n",
        "          assert self.config.agent_id not in c.train_locks\n",
        "    # Add locks.\n",
        "    paths = self.paths[self.config.agent_id]\n",
        "    for p in paths:\n",
        "      for c in p.components:\n",
        "        c.train_locks.add(self.config.agent_id)\n",
        "\n",
        "  def rm_train_locks(self):\n",
        "    # Remove locks.\n",
        "    paths = self.paths[self.config.agent_id]\n",
        "    for p in paths:\n",
        "      for c in p.components:\n",
        "        if self.config.agent_id in c.train_locks:\n",
        "          c.train_locks.remove(self.config.agent_id)\n",
        "    # Check.\n",
        "    for ps in self.paths.values():\n",
        "      for p in ps:\n",
        "        for c in p.components:\n",
        "          assert self.config.agent_id not in c.train_locks\n",
        "\n",
        "  def start_cycle(self):\n",
        "    self.rm_train_locks()\n",
        "\n",
        "  def end_cycle(self):\n",
        "    # Keep only best one.\n",
        "    best_path = self.get_best_path()\n",
        "    assert best_path is not None\n",
        "    best_path.metrics[\"num_cycles\"] = best_path.metrics.get(\"num_cycles\", 0) + 1\n",
        "    self.paths[self.config.agent_id] = [best_path]\n",
        "\n",
        "    self.add_train_locks()\n",
        "    self.garbage_collect_paths()\n",
        "\n",
        "  def garbage_collect_paths(self):\n",
        "    # Store history before dropping references to unused paths to trigger\n",
        "    # garbage collection of components and parameters.\n",
        "    self.paths_df = self.paths_df.append(\n",
        "        paths_to_df(Path.paths), ignore_index=True\n",
        "        ).query(f'agent_id==\"{self.config.agent_id}\" and id\u003e{Path.last_saved}')\n",
        "    self.comps_df = self.comps_df.append(\n",
        "        components_to_df(Path.paths), ignore_index=True\n",
        "        ).query(f'agent_id==\"{self.config.agent_id}\" and id\u003e{Component.last_saved}'\n",
        "        ).drop_duplicates()\n",
        "\n",
        "    # Drop unused paths generated in this task iteration for garbage collection.\n",
        "    Path.paths = []\n",
        "    # Simplify ancestor tree to contain only live paths.\n",
        "    live_paths_ids = [p.id for paths in self.paths.values() for p in paths]\n",
        "    # Notice that the simplification is done also for paths of other tasks,\n",
        "    # since they may be pointing to a path of this task that was just pruned.\n",
        "    for path in [path for paths in self.paths.values() for path in paths]:\n",
        "      ancestor = path.parent\n",
        "      if ancestor is None:\n",
        "        continue\n",
        "      while True:\n",
        "        if ancestor.id in live_paths_ids:\n",
        "          path.parent = ancestor\n",
        "          break\n",
        "        ancestor = ancestor.parent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4EgQHbawpNcS"
      },
      "outputs": [],
      "source": [
        "pd.set_option(\"display.expand_frame_repr\", False)\n",
        "pd.set_option(\"display.max_columns\", 100)\n",
        "pd.set_option(\"display.max_rows\", 100)\n",
        "\n",
        "def pop_to_df(pop):\n",
        "  return paths_to_df([p for paths in pop.paths.values() for p in paths])\n",
        "\n",
        "def paths_to_df(paths):\n",
        "  # Collect all metrics names.\n",
        "  metrics_keys = set()\n",
        "  hparams_keys = set()\n",
        "  for path in paths:\n",
        "    path.score()  # Update scores.\n",
        "    metrics_keys.update(path.metrics)\n",
        "    hparams_keys.update(path.hparams)\n",
        "\n",
        "  data = defaultdict(list)\n",
        "  for path in paths:\n",
        "    data[\"agent_id\"].append(path.agent_id)\n",
        "    data[\"task_name\"].append(path.task_name)\n",
        "    data[\"id\"].append(path.id)\n",
        "    data[\"parent_id\"].append(path.parent.id if path.parent else -1)\n",
        "    data[\"parent_agent_id\"].append(path.parent.agent_id if path.parent else None)\n",
        "    data[\"components\"].append(\",\".join(\n",
        "        [f\"{c.agent_id}:{c.id}\" for c in path.components]))\n",
        "    for k in hparams_keys:\n",
        "      data[f\"hparams.{k}\"].append(path.hparams[k] if k in path.hparams else None)\n",
        "    for k in metrics_keys:\n",
        "      data[f\"metrics.{k}\"].append(path.metrics[k] if k in path.metrics else None)\n",
        "  return pd.DataFrame(data)\n",
        "\n",
        "def components_to_df(paths):\n",
        "  # Collect all components.\n",
        "  comps = set()\n",
        "  for p in paths:\n",
        "    comps.update(p.components)\n",
        "\n",
        "  data = defaultdict(list)\n",
        "  for c in comps:\n",
        "    data[\"id\"].append(c.id)\n",
        "    data[\"name\"].append(c.name)\n",
        "    data[\"agent_id\"].append(c.agent_id)\n",
        "    data[\"num_params\"].append(c.get_num_params())\n",
        "  return pd.DataFrame(data)\n",
        "\n",
        "def print_df_segments(df, segment_length:int = 10):\n",
        "  tot_length = df.shape[0]\n",
        "  # Pad column title with spaces to keep alignment across segments.\n",
        "  def prepend_spaces(original_str, pad_to_len):\n",
        "    return \" \" * (pad_to_len-len(original_str)) + original_str\n",
        "  pad_to_len = max([len(tn) for tn in set(df[\"agent_id\"].to_list())])+1\n",
        "  df = df.rename(columns={\n",
        "    \"agent_id\": prepend_spaces(\"agent_id\", pad_to_len),\n",
        "    \"task_name\": prepend_spaces(\"task_name\", pad_to_len),\n",
        "    \"parent_agent_id\": prepend_spaces(\"parent_agent_id\", pad_to_len),\n",
        "    })\n",
        "  for x in range(0, tot_length, segment_length):\n",
        "    print(df[x:min(x+segment_length, tot_length)])\n",
        "\n",
        "def df_leaderboard(df):\n",
        "  # Place columns on the left for readability.\n",
        "  all_keys = sorted(df.columns.tolist())\n",
        "  first_keys = [\"agent_id\", \"task_name\",\n",
        "                \"metrics.test_quality\", \"metrics.score\", \"metrics.quality\",\n",
        "                \"metrics.accounted_params\", \"metrics.flops\",\n",
        "                \"id\", \"parent_id\", \"parent_agent_id\"]\n",
        "  first_keys = [k for k in first_keys if k in all_keys]\n",
        "  sorted_keys = first_keys + [k for k in all_keys if k not in first_keys]\n",
        "  # Filter mu function parameters.\n",
        "  sorted_keys = [k for k in sorted_keys if \"_mu_|\" not in k]\n",
        "  df = df[sorted_keys]\n",
        "  if \"metrics.score\" in df:\n",
        "    df = df.sort_values([\"agent_id\", \"metrics.score\"], ascending=[True, False],\n",
        "                        ignore_index=True)\n",
        "  else:\n",
        "    df = df.sort_values(\"agent_id\", ignore_index=True)\n",
        "  print_df_segments(df)\n",
        "  for k in [\"metrics.score\", \"metrics.quality\", \"metrics.test_quality\",\n",
        "            \"metrics.accounted_params\", \"metrics.flops\"]:\n",
        "    if k in df:\n",
        "      print(f\"Avg {k}: {df[k].mean():.6f}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xmYkTEY-PBgh"
      },
      "outputs": [],
      "source": [
        "def prp(path):\n",
        "  rtn = []\n",
        "  if VERBOSE:\n",
        "    rtn.append(str(path))\n",
        "    for c in path.components:\n",
        "      rtn.append(str(c))\n",
        "  else:\n",
        "    rtn.append(str(path.id))\n",
        "  return \"\\n\".join(rtn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rDU0butVlu2f"
      },
      "source": [
        "# Checkpointing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "104ReuGLspoJ"
      },
      "outputs": [],
      "source": [
        "def df_write_to_csv(df, dir_path, df_name):\n",
        "  filename_df = os.path.join(dir_path, f\"{df_name}.csv\")\n",
        "  with gfile.GFile(filename_df, \"w\") as outfile:\n",
        "    df.to_csv(outfile, index=False)\n",
        "\n",
        "def df_read_from_csv(dir_path, df_name):\n",
        "  filename_df = os.path.join(dir_path, f\"{df_name}.csv\")\n",
        "  with gfile.GFile(filename_df, \"r\") as infile:\n",
        "    df = pd.read_csv(infile)\n",
        "  # Pandas read_csv() reads empty stings as NaNs. Set NaNs to empty strings in\n",
        "  # columns with type strings/object.\n",
        "  for c in df.columns:\n",
        "    if df[c].dtype == np.object_:\n",
        "        df[c].fillna(\"\", inplace=True)\n",
        "  return df\n",
        "\n",
        "def get_comps_params_to_save(pop: Population):\n",
        "  \"\"\"Returns a dictionary containing the parameters of the used components.\"\"\"\n",
        "  comps_params = {}\n",
        "  # All components generated by this agent.\n",
        "  all_comps = set(\n",
        "      [c for p in pop.paths[pop.config.agent_id] for c in p.components if c.agent_id == pop.config.agent_id])\n",
        "  # Check that there are not duplicate ids.\n",
        "  assert len(all_comps) == len(set([c.id for c in all_comps])), (\n",
        "      [f\"{c.name}:{c.agent_id}:{c.id}\" for c in all_comps])\n",
        "  for c in all_comps:\n",
        "    if c.id \u003c= Component.last_saved:\n",
        "      continue\n",
        "    assert c.agent_id == pop.config.agent_id\n",
        "    c_id_string = f\"{c.name}:{c.agent_id}:{c.id}\"\n",
        "    comps_params[c_id_string] = c.params\n",
        "    if c.opt_state is not None:\n",
        "      comps_params[f\"opt_state:{c_id_string}\"] = c.opt_state\n",
        "  return comps_params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z_AmkV7hLeQp"
      },
      "outputs": [],
      "source": [
        "def latest_checkpoint(ckpt_dir: str, prefix: str = \"checkpoint_\"):\n",
        "  \"\"\"Returns the latest checkpoint under the dir pattern or None if missing.\"\"\"\n",
        "  ckpt_dir = os.fspath(ckpt_dir)\n",
        "  glob_path = os.path.join(ckpt_dir, f\"{prefix}*\")\n",
        "  checkpoint_files = flax_checkpoints.natural_sort(gfile.glob(glob_path))\n",
        "  checkpoint_files = [f for f in checkpoint_files if not f.endswith(\"_tmp\")]\n",
        "  return checkpoint_files[-1] if checkpoint_files else None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M_RTFjoOwfpM"
      },
      "outputs": [],
      "source": [
        "LAST_CHECKPOINT_TIME = time.time()\n",
        "\n",
        "def save_checkpoint(\n",
        "    ckpt_dir: str,\n",
        "    comps_params,\n",
        "    cycle_id: int,\n",
        "    generation_id: int):\n",
        "  print(\"SAVING\", cycle_id, generation_id, comps_params.keys())\n",
        "  # Write checkpoint.\n",
        "  flax_checkpoints.save_checkpoint(\n",
        "      ckpt_dir=ckpt_dir,\n",
        "      target=comps_params,\n",
        "      step=generation_id,\n",
        "      prefix=f\"checkpoint_{cycle_id}_\",\n",
        "      overwrite=True)\n",
        "  # Update time of last checkpoint save.\n",
        "  global LAST_CHECKPOINT_TIME\n",
        "  LAST_CHECKPOINT_TIME = time.time()\n",
        "  # Delete intermediate checkpoint directories.\n",
        "  if generation_id == 0:\n",
        "    intermediate_ckpt_dirs = gfile.glob(\n",
        "        os.path.join(os.path.dirname(ckpt_dir), \"state_*_[^0]*\"))\n",
        "    for d in intermediate_ckpt_dirs:\n",
        "      print(\"Deleting intermediate checkpoint:\", d)\n",
        "      gfile.rmtree(d)\n",
        "\n",
        "def save_state(\n",
        "    pop: Population,\n",
        "    cycle_id: int,\n",
        "    generation_id: int,\n",
        "    config: FrozenConfigDict):\n",
        "  \"\"\"Save checkpoint and data needed to resume exp.\"\"\"\n",
        "  intermediate = (generation_id != 0)\n",
        "  write_threads = []\n",
        "  # Skip intermediate state save if last state was written recently.\n",
        "  if intermediate and (\n",
        "      (time.time() - LAST_CHECKPOINT_TIME) \u003c SKIP_INTERMEDIATE_STATE_SECS):\n",
        "    print(\"Skip checkpointing, seconds since last save:\",\n",
        "          f\"{time.time() - LAST_CHECKPOINT_TIME:.0f}\")\n",
        "    return write_threads\n",
        "\n",
        "  # Save data needed to resume exp.\n",
        "  write_start = time.time()\n",
        "  df_leaderboard(pop_to_df(pop))\n",
        "  pop.garbage_collect_paths()\n",
        "  state_dir = os.path.join(\n",
        "      config.agent_dir, f\"state_{cycle_id}_{generation_id}\")\n",
        "  gfile.makedirs(state_dir)\n",
        "\n",
        "  if latest_checkpoint(state_dir):\n",
        "    assert False, f\"checkpoint already present in forlder: {state_dir}\"\n",
        "  print(\"WRITING CHECKPOINT:\", cycle_id, generation_id)\n",
        "\n",
        "  # Write state in background threads.\n",
        "  write_threads = []\n",
        "\n",
        "  write_threads.append(Thread(\n",
        "      target=df_write_to_csv,\n",
        "      args=(paths_to_df([p for p in pop.paths[config.agent_id]]),\n",
        "            state_dir,\n",
        "            \"published\")))\n",
        "  write_threads[-1].start()\n",
        "\n",
        "  write_threads.append(Thread(\n",
        "      target=df_write_to_csv,\n",
        "      args=(paths_to_df([p for paths in pop.paths.values() for p in paths]),\n",
        "            state_dir,\n",
        "            \"population\")))\n",
        "  write_threads[-1].start()\n",
        "\n",
        "  write_threads.append(Thread(\n",
        "      target=df_write_to_csv,\n",
        "      args=(pop.paths_df, state_dir, \"paths\")))\n",
        "  write_threads[-1].start()\n",
        "\n",
        "  write_threads.append(Thread(\n",
        "      target=df_write_to_csv,\n",
        "      args=(pop.comps_df, state_dir, \"components\")))\n",
        "  write_threads[-1].start()\n",
        "\n",
        "  write_threads.append(Thread(\n",
        "      target=json.dump,\n",
        "      args=(config.as_configdict().to_dict(),\n",
        "            gfile.GFile(os.path.join(state_dir, \"config.json\"), \"w\")),\n",
        "      kwargs=dict(indent=2)))\n",
        "  write_threads[-1].start()\n",
        "\n",
        "  write_threads.append(Thread(\n",
        "      target=save_checkpoint,\n",
        "      args=(state_dir, get_comps_params_to_save(pop), cycle_id, generation_id)))\n",
        "  write_threads[-1].start()\n",
        "\n",
        "  # Update last saved.\n",
        "  if not intermediate:\n",
        "    Path.last_saved = pop.paths_df.id.max()\n",
        "    Component.last_saved = pop.comps_df.id.max()\n",
        "\n",
        "  print(f\"WRITE START TIME: {time.time() - write_start:.2f} s\")\n",
        "  return write_threads"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7yWdy7DskBph"
      },
      "outputs": [],
      "source": [
        "def load_paths(\n",
        "    pop: Population,\n",
        "    state_dir: str,\n",
        "    all_agents_dirs: str,\n",
        "    ):\n",
        "  if state_dir:\n",
        "    state_dir = state_dir.rstrip(\"/\")\n",
        "\n",
        "  load_start = time.time()\n",
        "\n",
        "  # Load system state info.\n",
        "  population_df = pd.DataFrame()\n",
        "  skip_agent_dir = None\n",
        "  if state_dir:\n",
        "    # Load active agent state, possibly intermediate.\n",
        "    population_df = population_df.append(\n",
        "        df_read_from_csv(state_dir, \"published\"))\n",
        "    skip_agent_dir = os.path.dirname(state_dir)\n",
        "  for agent_dir in gfile.glob(all_agents_dirs):\n",
        "    if agent_dir == skip_agent_dir:\n",
        "      continue\n",
        "    agent_checkpoint = latest_checkpoint(os.path.join(agent_dir, \"state_*_0/\"))\n",
        "    if agent_checkpoint is None:\n",
        "      continue\n",
        "    population_df = population_df.append(\n",
        "        df_read_from_csv(os.path.dirname(agent_checkpoint), \"published\"))\n",
        "\n",
        "  # Load parameters from sharded system checkpoint.\n",
        "  loaded_params = {}  # Dictionary to accumlate loaded parameters.\n",
        "  lock = Lock()\n",
        "  duplicate_keys = set()\n",
        "  def append_loaded_params(add_chkp_dir: str):\n",
        "    # Skip folders without a completed checkpoint.\n",
        "    if latest_checkpoint(add_chkp_dir) is None:\n",
        "      return\n",
        "    lp_add = flax_checkpoints.restore_checkpoint(\n",
        "        ckpt_dir=add_chkp_dir,\n",
        "        target=None)\n",
        "    if lp_add:\n",
        "      lock.acquire()\n",
        "      print(\"LOADED COMPONENTS\", add_chkp_dir, lp_add.keys())\n",
        "      duplicate_keys.update(loaded_params.keys() \u0026 lp_add.keys())\n",
        "      loaded_params.update(lp_add)\n",
        "      lock.release()\n",
        "  all_state_dirs = []\n",
        "  if state_dir:\n",
        "    # Include active agent state, possibly intermediate.\n",
        "    all_state_dirs.append(state_dir)\n",
        "    all_state_dirs.extend(\n",
        "        gfile.glob(os.path.join(os.path.dirname(state_dir), \"state_*_0\"))\n",
        "    )\n",
        "  for agent_dir in gfile.glob(all_agents_dirs):\n",
        "    all_state_dirs.extend(\n",
        "        gfile.glob(os.path.join(agent_dir, \"state_*_0\"))\n",
        "    )\n",
        "  threads = []\n",
        "  for s_dir in set(all_state_dirs):\n",
        "    threads.append(Thread(target=append_loaded_params, args=(s_dir,)))\n",
        "    threads[-1].start()\n",
        "  for t in threads:\n",
        "    t.join()\n",
        "  assert not duplicate_keys, duplicate_keys\n",
        "\n",
        "  print(f\"LOAD TIME: {time.time() - load_start:.2f} s\")\n",
        "\n",
        "  frozen_params = flax.core.freeze(loaded_params)\n",
        "  sid_2_comp = {}\n",
        "  for k in frozen_params.keys():\n",
        "    if k.startswith(\"opt_state:\"):\n",
        "      continue\n",
        "    assert len(k.split(\":\")) == 3, k\n",
        "    name, agent_id, id = k.split(\":\")\n",
        "    if \"opt_state:\" + k in frozen_params.keys():\n",
        "      opt_state = frozen_params[\"opt_state:\" + k]\n",
        "    else:\n",
        "      opt_state = None\n",
        "    c = Component(\n",
        "        name=name, agent_id=agent_id,\n",
        "        params=frozen_params[k], train_locks=[], opt_state=opt_state)\n",
        "    c.id = int(id)\n",
        "    source_id = f\"{agent_id}:{id}\"\n",
        "    assert source_id not in sid_2_comp, source_id\n",
        "    sid_2_comp[source_id] = c\n",
        "  # For parent assignemt.\n",
        "  sid_2_path = {}\n",
        "  path_2_parent_sid = {}\n",
        "  for index, row in population_df.iterrows():\n",
        "    comps_sids = row[\"components\"].split(\",\")\n",
        "    comps = []\n",
        "    for sid in comps_sids:\n",
        "      comps.append(sid_2_comp[sid])\n",
        "    task_name = row[\"task_name\"]\n",
        "    agent_id = row[\"agent_id\"]\n",
        "\n",
        "    # Retrieve hparams and metrics.\n",
        "    hparams = {}\n",
        "    metrics = {}\n",
        "    for k in row.keys():\n",
        "      if type(row[k]) is float and math.isnan(row[k]):\n",
        "        continue\n",
        "      if k.startswith(\"hparams.\"):\n",
        "        hparams[k[len(\"hparams.\"):]] = row[k]\n",
        "      if k.startswith(\"metrics.\"):\n",
        "        metrics[k[len(\"metrics.\"):]] = row[k]\n",
        "    # Fix adapter_layers string.\n",
        "    if \"adapter_layers\" in hparams and type(hparams[\"adapter_layers\"]) is float:\n",
        "      hparams[\"adapter_layers\"] = str(int(hparams[\"adapter_layers\"]))\n",
        "    # Create path.\n",
        "    path = Path(\n",
        "        hparams=hparams,\n",
        "        components=comps,\n",
        "        parent=None,\n",
        "        agent_id=agent_id,\n",
        "        task_name=task_name,\n",
        "        )\n",
        "    path.metrics = metrics\n",
        "    path.id = int(row[\"id\"])\n",
        "    # Add train locks.\n",
        "    for c in path.components:\n",
        "      c.train_locks.add(agent_id)\n",
        "    pop.paths[agent_id].append(path)\n",
        "    path_sid = f\"{path.agent_id}/{path.id}\"\n",
        "    assert path_sid not in sid_2_path\n",
        "    sid_2_path[path_sid] = path\n",
        "    if row[\"parent_id\"] \u003e= 0:\n",
        "      parent_sid = f'{row[\"parent_agent_id\"]}/{row[\"parent_id\"]}'\n",
        "      path_2_parent_sid[path] = parent_sid\n",
        "\n",
        "  # Set parents.\n",
        "  for path, parent_sid in path_2_parent_sid.items():\n",
        "    if parent_sid not in sid_2_path:\n",
        "      # This can happen if parent is retired by a parallel agent.\n",
        "      # In this case fall back to root model.\n",
        "      for k in sid_2_path.keys():\n",
        "        if \"root_model\"in k:\n",
        "          parent_sid = k\n",
        "      print(f\"{path.agent_id}:{path.id} orphaned, fallback: {parent_sid}\")\n",
        "    path.parent = sid_2_path[parent_sid]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_0AyzNaTl002"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m6vSjSIvNPq4"
      },
      "outputs": [],
      "source": [
        "@partial(jax.jit, static_argnames=\"model\")\n",
        "def eval_step(params, images, labels, model):\n",
        "  logits = model.apply({\"params\": params}, images,\n",
        "                       train=False)  # Disable dropout.\n",
        "  # Avg accuracy on the batch.\n",
        "  return (logits.argmax(axis=-1) == labels).mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_m2xl8XR7cWy"
      },
      "outputs": [],
      "source": [
        "@partial(jax.jit, static_argnames=[\"model\", \"optimizer\"], donate_argnums=[0, 2])\n",
        "def train_step(params, fixed_params, opt_state, images, labels, model, optimizer):\n",
        "  def loss_fn(params, fixed_params, images, labels):\n",
        "    logits = model.apply(\n",
        "        {\"params\": format_params(params, fixed_params)}, images,\n",
        "        train=False)  # Disable dropout.\n",
        "    labels = jax.nn.one_hot(labels, logits.shape[-1])\n",
        "    return -jnp.mean(jnp.sum(labels * nn.log_softmax(logits), axis=-1))\n",
        "  grads = jax.grad(loss_fn)(params, fixed_params, images, labels)\n",
        "  updates, opt_state = optimizer.update(grads, opt_state, params=params)\n",
        "  params = optax.apply_updates(params, updates)\n",
        "  return params, opt_state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lRjic_IpGYJU"
      },
      "outputs": [],
      "source": [
        "LOOP_START = time.time()\n",
        "\n",
        "def train_loop(paths, ds_train, ds_validation, devices, config):\n",
        "  global LOOP_START\n",
        "  timing = {\"start_time\": time.time(),\n",
        "            \"start_time_loop\": LOOP_START}\n",
        "  task = paths[0].task\n",
        "  # The following values should be shared by all paths in this generation batch.\n",
        "  for path in paths:\n",
        "    assert task.name == path.task_name\n",
        "    assert paths[0].hparams[\"ds_image_size\"] == path.hparams[\"ds_image_size\"]\n",
        "\n",
        "  gc.collect()\n",
        "\n",
        "  # Compile.\n",
        "  compile_train_batches_arr = jax.device_put_replicated(\n",
        "      get_sample_batch(\n",
        "        paths[0].hparams[\"ds_image_size\"],\n",
        "        task.train_batch_size),\n",
        "      devices)\n",
        "  compile_eval_batches_arr = jax.device_put_replicated(\n",
        "      get_sample_batch(\n",
        "          paths[0].hparams[\"ds_image_size\"],\n",
        "          task.validation_batch_size),\n",
        "      devices)\n",
        "\n",
        "  for p_id, path in enumerate(paths):\n",
        "    if VERBOSE:\n",
        "      print(\"Parent\")\n",
        "      print(prp(path.parent))\n",
        "      print(prp(path))\n",
        "    path.device_id = p_id % len(devices)\n",
        "    path.device = devices[path.device_id]\n",
        "    path.optimizer = path.get_optimizer()\n",
        "    path.optimizer_init_fn = jax.jit(path.optimizer.init, device=path.device)\n",
        "    path.best_params_local = None\n",
        "    path.best_opt_state_local = None\n",
        "    path.best_quality = None\n",
        "    path.best_score = path.parent.score() if path.task_name == path.parent.task_name else -np.inf\n",
        "    path.evals = []\n",
        "\n",
        "    # Launch parallel compilation of eval and train step functions.\n",
        "    params_local = path.get_trainable_params()\n",
        "    path.compile_params_device = jax.device_put(params_local, path.device)\n",
        "    path.compile_fixed_params_device = jax.device_put(\n",
        "        path.get_fixed_params(),\n",
        "        path.device)\n",
        "    path.compile_train = Thread(\n",
        "        target=train_step,\n",
        "        args=(path.compile_params_device,\n",
        "              path.compile_fixed_params_device,\n",
        "              path.optimizer_init_fn(params_local),\n",
        "              compile_train_batches_arr[\"image\"][path.device_id],\n",
        "              compile_train_batches_arr[\"label\"][path.device_id],\n",
        "              path.model,\n",
        "              path.optimizer))\n",
        "    path.compile_eval = Thread(\n",
        "        target=eval_step,\n",
        "        args=(format_params(\n",
        "                  path.compile_params_device,\n",
        "                  path.compile_fixed_params_device),\n",
        "              compile_eval_batches_arr[\"image\"][path.device_id],\n",
        "              compile_eval_batches_arr[\"label\"][path.device_id],\n",
        "              path.model))\n",
        "    path.compile_eval.start()\n",
        "\n",
        "  for path in paths:\n",
        "    path.compile_eval.join()\n",
        "    del path.compile_eval\n",
        "    timing[\"end_compile_eval\"] = time.time()\n",
        "    path.compile_train.start()\n",
        "  del compile_eval_batches_arr\n",
        "\n",
        "  for path in paths:\n",
        "    path.compile_train.join()\n",
        "    del path.compile_train\n",
        "    del path.compile_params_device\n",
        "    del path.compile_fixed_params_device\n",
        "    timing[\"end_compile\"] = time.time()\n",
        "  del compile_train_batches_arr\n",
        "\n",
        "  gc.collect()\n",
        "\n",
        "  # Parameter tranfer.\n",
        "  for path in paths:\n",
        "    path.params_device = jax.device_put(\n",
        "        path.get_trainable_params(),\n",
        "        path.device)\n",
        "    path.fixed_params_device = jax.device_put(\n",
        "        path.get_fixed_params(),\n",
        "        path.device)\n",
        "    path.opt_state_device = path.optimizer_init_fn(path.params_device)\n",
        "    # Set opt state.\n",
        "    for c in path.components:\n",
        "      if c.is_trainable():\n",
        "        assert c.name in path.opt_state_device[1][0].trace.keys()\n",
        "        if c.opt_state is not None:\n",
        "          path.opt_state_device = (\n",
        "              path.opt_state_device[0],\n",
        "              (optax.TraceState(\n",
        "                  trace=path.opt_state_device[1][0].trace.copy(\n",
        "                      {c.name: jax.device_put(c.opt_state,\n",
        "                                              path.device)})),\n",
        "               path.opt_state_device[1][1]\n",
        "               )\n",
        "          )\n",
        "\n",
        "  iter_ds_validation = iter(ds_validation)\n",
        "  # TRAIN\n",
        "  for t_step, train_batch in zip(\n",
        "      range(config.num_validations_per_path_training\n",
        "            * task.num_train_batches_between_validations),\n",
        "      ds_train,\n",
        "  ):\n",
        "    train_batch_arr = jax.device_put_replicated(train_batch, devices)\n",
        "    for p_id, path in enumerate(paths):\n",
        "      if t_step == 0:\n",
        "        timing[\"end_prep\"] = time.time()\n",
        "        t_step_0_time = time.time()\n",
        "      path.params_device, path.opt_state_device = train_step(\n",
        "          path.params_device,\n",
        "          path.fixed_params_device,\n",
        "          path.opt_state_device,\n",
        "          train_batch_arr[\"image\"][path.device_id],\n",
        "          train_batch_arr[\"label\"][path.device_id],\n",
        "          path.model,\n",
        "          path.optimizer)\n",
        "      if t_step == 0 and time.time() - t_step_0_time \u003e 1:\n",
        "        print(f\"WARNING: First train step took: {time.time()-t_step_0_time:.2f} s\")\n",
        "    del train_batch, train_batch_arr\n",
        "\n",
        "    # EVAL\n",
        "    if (t_step+1) % task.num_train_batches_between_validations == 0:\n",
        "      first_eval = ((t_step+1) == task.num_train_batches_between_validations)\n",
        "      if first_eval:\n",
        "        timing[\"start_eval\"] = time.time()\n",
        "      for path in paths:\n",
        "        path.accs = []\n",
        "      for e_step, eval_batch in zip(\n",
        "          range(task.num_validation_batches),\n",
        "          iter_ds_validation,\n",
        "          ):\n",
        "        eval_batch_arr = jax.device_put_replicated(eval_batch, devices)\n",
        "        for p_id, path in enumerate(paths):\n",
        "          if first_eval and e_step == 0:\n",
        "            e_step_0_time = time.time()\n",
        "          path.accs.append(\n",
        "              eval_step(\n",
        "                  format_params(path.params_device, path.fixed_params_device),\n",
        "                  eval_batch_arr[\"image\"][path.device_id],\n",
        "                  eval_batch_arr[\"label\"][path.device_id],\n",
        "                  path.model))\n",
        "          if first_eval and e_step == 0 and time.time() - e_step_0_time \u003e 1:\n",
        "            print(f\"WARNING: First eval step took: {time.time()-e_step_0_time:.2f} s\")\n",
        "      del eval_batch, eval_batch_arr\n",
        "\n",
        "      # Get params of best models.\n",
        "      qs = []\n",
        "      eval_idx = (t_step+1) // task.num_train_batches_between_validations\n",
        "      for path in paths:\n",
        "        quality = np.mean(path.accs)\n",
        "        del path.accs\n",
        "        qs.append(f\"{quality:.4f}\")\n",
        "        path.evals.append(quality)\n",
        "        # Set quality in metrics for current score computation.\n",
        "        path.metrics[\"quality\"] = quality\n",
        "        path_score = path.score()\n",
        "        if path_score \u003e= path.best_score:\n",
        "          path.best_params_local = jax.device_get(path.params_device)\n",
        "          path.best_opt_state_local = jax.device_get(path.opt_state_device[1][0].trace)\n",
        "          path.best_score = path_score\n",
        "          path.best_quality = quality\n",
        "          qs[-1] += \"*\"\n",
        "      train_time = time.time() - timing[\"end_compile\"]\n",
        "      avg_path_time = (train_time / eval_idx) / len(paths)\n",
        "      print((\"\\t\".join(qs) + f\"\\t\u003c Eval {eval_idx}\").expandtabs(8),\n",
        "            f\"tot:{train_time:.1f}s\", f\"avg/path:{avg_path_time:.1f}s\")\n",
        "\n",
        "      if first_eval:\n",
        "        timing[\"end_eval\"] = time.time()\n",
        "\n",
        "  for path in paths:\n",
        "    del path.params_device\n",
        "    del path.fixed_params_device\n",
        "    del path.opt_state_device\n",
        "    del path.optimizer\n",
        "    del path.optimizer_init_fn\n",
        "  gc.collect()\n",
        "\n",
        "  timing[\"end_train\"] = time.time()\n",
        "\n",
        "  loop_time = timing[\"start_time\"] - LOOP_START\n",
        "  compile_time = timing[\"end_compile\"] - timing[\"start_time\"]\n",
        "  compile_eval_time = timing[\"end_compile_eval\"] - timing[\"start_time\"]\n",
        "  compile_train_time = timing[\"end_compile\"] - timing[\"end_compile_eval\"]\n",
        "  prep_time = timing[\"end_prep\"] - timing[\"end_compile\"]\n",
        "  train_time = timing[\"end_train\"] - timing[\"end_prep\"]\n",
        "  eval_time = timing[\"end_eval\"] - timing[\"start_eval\"]\n",
        "  LOOP_START = time.time()\n",
        "\n",
        "  for path in paths:\n",
        "    path.metrics[\"loop_time\"] = loop_time\n",
        "    path.metrics[\"compile_time\"] = compile_time\n",
        "    path.metrics[\"prep_time\"] = prep_time\n",
        "    path.metrics[\"train_time\"] = train_time\n",
        "    path.metrics[\"eval_time\"] = eval_time\n",
        "    path.metrics[\"start_time\"] = timing[\"start_time\"]\n",
        "    path.metrics[\"start_time_loop\"] = timing[\"start_time_loop\"]\n",
        "    path.metrics[\"end_time\"] = time.time()\n",
        "    path.metrics[\"num_params\"] = get_num_params(path.get_all_params())\n",
        "    path.metrics[\"num_trainable_params\"] = get_num_params(path.get_trainable_params())\n",
        "    path.metrics[\"quality\"] = max(path.evals)\n",
        "    path.metrics[\"evals\"] = json.dumps([float(v) for v in path.evals])\n",
        "\n",
        "    if path.best_params_local:\n",
        "      path.metrics[\"improved\"] = True\n",
        "      path.update_trainable(path.best_params_local,\n",
        "                            path.best_opt_state_local)\n",
        "      assert path.best_quality == path.metrics[\"quality\"]\n",
        "      assert path.best_score == path.score()\n",
        "    else:\n",
        "      path.metrics[\"improved\"] = False\n",
        "      # Sampled path will be dropped if not improved, so skip paramter update.\n",
        "      assert path.best_params_local == None\n",
        "      assert path.best_opt_state_local == None\n",
        "      assert path.best_quality == None\n",
        "\n",
        "    del path.best_params_local\n",
        "    del path.best_opt_state_local\n",
        "    del path.best_score\n",
        "    del path.best_quality\n",
        "    del path.evals\n",
        "\n",
        "    if VERBOSE:\n",
        "      print(\"UPDATED:\")\n",
        "      print(prp(path))\n",
        "\n",
        "  pqs = []\n",
        "  qs = []\n",
        "  psc = []\n",
        "  sc = []\n",
        "  for path in paths:\n",
        "    if path.task_name == path.parent.task_name:\n",
        "      pqs.append(f\"{path.parent.metrics['quality']:.4f}\")\n",
        "      psc.append(f\"{path.parent.score():.4f}\")\n",
        "    else:\n",
        "      pqs.append(\"NEW\")\n",
        "      psc.append(\"NEW\")\n",
        "    qs.append(f\"{path.metrics['quality']:.4f}\")\n",
        "    sc.append(f\"{path.score():.4f}\")\n",
        "    if path.metrics[\"improved\"]:\n",
        "      sc[-1] += \"+\"\n",
        "\n",
        "  print((\"\\t\".join([f\"{path.parent.id}\" for path in paths]) +\n",
        "        \"\\t\u003c Parent id\").expandtabs(8))\n",
        "  print((\"\\t\".join([f\"{path.id}\" for path in paths]) +\n",
        "        \"\\t\u003c Path id\").expandtabs(8))\n",
        "  print((\"\\t\".join(pqs) + \"\\t\u003c Parent best quality\").expandtabs(8))\n",
        "  print((\"\\t\".join(qs) + \"\\t\u003c Path best quality\").expandtabs(8))\n",
        "  print((\"\\t\".join(psc) + \"\\t\u003c Parent score\").expandtabs(8))\n",
        "  print((\"\\t\".join(sc) + \"\\t\u003c Path score\").expandtabs(8))\n",
        "\n",
        "  print(\"time\\tINIT\\tCOMPevl\\tCOMPtrn\\tPREP\\tTRN+EVL\\t1stEVAL\".expandtabs(8))\n",
        "  print(f\"(s)\\t{loop_time:.1f}\\t{compile_eval_time:.1f}\\t{compile_train_time:.1f}\\t{prep_time:.1f}\\t{train_time:.1f}\\t{eval_time:.1f}\".expandtabs(8))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p-3gYmECRBOU"
      },
      "outputs": [],
      "source": [
        "# Run a full paths sampling iteration for a task.\n",
        "def agent_cycle(\n",
        "    task, devices, pop:Population, generation_id:int, cycle_id:int,\n",
        "    config:FrozenConfigDict):\n",
        "  num_devices = len(devices)\n",
        "  # Track best path.\n",
        "  best_path = pop.get_best_path()\n",
        "  num_gen_batches = math.ceil(config.num_samples_per_cycle/num_devices)\n",
        "  for _ in range(num_gen_batches):\n",
        "    if generation_id \u003e= num_gen_batches:\n",
        "      break\n",
        "    print(\"----\")\n",
        "    print(f\"GENERATION: [{generation_id+1}/{num_gen_batches}]\")\n",
        "    ds_hparams = pop.sample_ds_hparams()\n",
        "    ds_hparams[\"num_classes\"] = task.num_classes\n",
        "    ds_train = task.get_ds(\"train\", ds_hparams)\n",
        "    ds_validation = task.get_ds(\"validation\", ds_hparams)\n",
        "    paths = []\n",
        "    for i in range(num_devices):\n",
        "      paths.append(pop.sample_path(ds_hparams))\n",
        "    train_loop(paths, ds_train, ds_validation, devices, config)\n",
        "    for path in paths:\n",
        "      if path.metrics[\"improved\"]:\n",
        "        assert path not in pop.paths[config.agent_id]\n",
        "        pop.paths[config.agent_id].append(path)\n",
        "    pop.prune_population()\n",
        "    # Track best path.\n",
        "    curr_best_path = pop.get_best_path()\n",
        "    if curr_best_path != best_path:\n",
        "      if best_path:\n",
        "        assert curr_best_path.score() \u003e= best_path.score()\n",
        "      best_path = curr_best_path\n",
        "      best_path.metrics[\"new_best\"] = True\n",
        "      print(f\"Best id:{best_path.id}\",\n",
        "            f\"score:{best_path.score():.4f}\",\n",
        "            f\"quality:{best_path.metrics['quality']:.4f}\",\n",
        "            f\"gen:{generation_id}\",\n",
        "            f\"\\n{best_path.hparams}\")\n",
        "    generation_id += 1\n",
        "    if generation_id \u003c num_gen_batches:\n",
        "      save_state(pop, cycle_id, generation_id, config)\n",
        "  assert best_path in pop.paths[config.agent_id]\n",
        "  run_test_eval(best_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JVT8nwIWMAVf"
      },
      "outputs": [],
      "source": [
        "def has_test_quality(path):\n",
        "  return (\"test_quality\" in path.metrics\n",
        "          and not math.isnan(path.metrics[\"test_quality\"]))\n",
        "\n",
        "# Run final eval on test set.\n",
        "def run_test_eval(path, test_immutability=False):\n",
        "  # Skip if test_quality already computed and no immutability test required.\n",
        "  if not test_immutability and has_test_quality(path):\n",
        "    return\n",
        "  eval_st = time.time()\n",
        "  ds_test = path.task.get_ds(\"test\", path.hparams)\n",
        "  # Running on same device should allow to reuse the fn compiled for validation\n",
        "  # if batch size matches.\n",
        "  params = path.get_all_params()\n",
        "  if not hasattr(path, \"device\"):\n",
        "    path.device = random.choice(jax.local_devices())\n",
        "  params_device = jax.device_put(params_comps_to_model(params), path.device)\n",
        "  acc_sum = []\n",
        "  tot_num_samples = 0\n",
        "  # Warning: if repeat() is called on this dataset, then this loop never ends.\n",
        "  for batch in ds_test:\n",
        "    acc_avg = eval_step(\n",
        "        params_device,\n",
        "        batch[\"image\"],\n",
        "        batch[\"label\"],\n",
        "        path.model)\n",
        "    batch_size = batch[\"image\"].shape[0]\n",
        "    # Need to recompute sum because last batch can have different size to allow\n",
        "    # for exact eval on the test set.\n",
        "    acc_sum.append(acc_avg * batch_size)\n",
        "    tot_num_samples += batch_size\n",
        "  del params_device\n",
        "  acc_avg = np.sum(acc_sum) / tot_num_samples\n",
        "  # Assert test quality equivalence to test immutability.\n",
        "  if has_test_quality(path):\n",
        "    print(f\"Testing immutability of path {path.id} : {acc_avg} ~= {path.metrics['test_quality']}\")\n",
        "    assert test_immutability\n",
        "    if not np.isclose(path.metrics[\"test_quality\"], acc_avg):\n",
        "      print(\"WARNING IMMUTABILITY TEST FAILED, delta:\",\n",
        "            path.metrics[\"test_quality\"]-acc_avg)\n",
        "  path.metrics[\"test_quality\"] = acc_avg\n",
        "  print(f\"TEST EVAL TIME: {time.time() - eval_st:.2f} s\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZubWF7rhmHju"
      },
      "source": [
        "# Main"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wzO_0X5-S0yj"
      },
      "outputs": [],
      "source": [
        "# Main loop over tasks.\n",
        "\n",
        "print(\"Devices type:\", jax.local_devices()[0].device_kind)\n",
        "print(\"Devices count:\", len(jax.local_devices()))\n",
        "\n",
        "agent = globals()[AGENT](\n",
        "    task_name=TASK_NAME,\n",
        "    experiment_dir=os.path.join(\n",
        "        EXPERIMENTS_ROOT_DIR,\n",
        "        EXPERIMENT_NAME),\n",
        "    auto_tune=AUTO_TUNE,\n",
        "    scale_factor=SCALE_FACTOR,\n",
        ")\n",
        "agent.run()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "name": "mu3Net: Multiagent Multitask Mutant Net",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
